Google का Tunix सिस्टम उस मुख्य बाधा (choke point) को दूर करता है जिसने बड़े पैमाने पर एजेंटिक रिइन्फोर्समेंट लर्निंग (RL) को TPUs का कुशलतापूर्वक उपयोग करने से रोक रखा था। इंटरैक्शन डेटा जेनरेट करने के काम को पॉलिसी अपडेट करने के काम से अलग करके, Tunix TPU उपयोग को सिंगल-डिजिट प्रतिशत से बढ़ाकर लगभग पूर्ण क्षमता तक ले जाता है, जिससे कंप्यूट बर्बादी में भारी कमी आती है।

एजेंटिक RL में बॉटलनेक

एजेंटिक RL, अधिक परिचित "नेक्स्ट-टोकन" लैंग्वेज-मॉडल ट्रेनिंग से अलग है। एक एजेंट को API कॉल भेजनी पड़ती है, कोड निष्पादित करना पड़ता है, या एक सिम्युलेटेड वातावरण में कदम बढ़ाना पड़ता है, और फिर परिणाम पर प्रतिक्रिया देनी पड़ती है। इसलिए ट्रेनिंग लूप सिंक्रोनस (synchronous) होता है: मॉडल एक एक्शन लेता है, वातावरण चलता है, परिणाम वापस आता है, और उसके बाद ही मॉडल को ग्रेडिएंट अपडेट प्राप्त होता है। जब एक सिंगल एनवायरनमेंट स्टेप में कई सेकंड लगते हैं, तो महंगा TPU हार्डवेयर खाली बैठा रहता है, और रिपोर्ट किया गया उपयोग 10% से नीचे गिर सकता है। यह अक्षमता सीधे तौर पर उच्च क्लाउड बिल और धीमी रिसर्च साइकिल में बदल जाती है।

Tunix का डिकपल्ड आर्किटेक्चर

Tunix दो चरणों—ट्राजेक्टरी जनरेशन और पॉलिसी ऑप्टिमाइज़ेशन—को अलग-अलग हार्डवेयर पूल्स पर लाकर इस समस्या का समाधान करता है।

  • Asynchronous actors सस्ते CPUs या GPUs पर चलते हैं। प्रत्येक एक्टर अपने सौंपे गए वातावरण के साथ लगातार इंटरैक्ट करता है, एक्शन और ऑब्जर्वेशन रिकॉर्ड करता है, और परिणामी ट्राजेक्टरीज़ को एक साझा स्टोर में स्ट्रीम करता है।
  • Continuous learners समर्पित TPU Pods पर होते हैं। लर्नर सेंट्रल बफर से बैच निकालता है और किसी भी सिंगल एक्टर के रोलआउट पूरा होने का इंतज़ार किए बिना ग्रेडिएंट अपडेट करता है।
  • High-throughput buffer बीच में स्थित होता है, जो ट्राजेक्टरीज़ के लिए स्टेजिंग एरिया के रूप में कार्य करता है। क्योंकि लर्नर उतनी ही तेज़ी से पढ़ सकता है जितनी तेज़ी से बफर डेटा सप्लाई कर सकता है, इसलिए TPU कभी रुकता नहीं है।

इसका कुल प्रभाव एक ऐसा ट्रेनिंग पाइपलाइन है जहाँ TPUs लगभग हर समय व्यस्त रहते हैं, जिससे उपयोग (utilization) 100% की ओर बढ़ता है।

तकनीकी बाधाएं और Tunix उन्हें कैसे दूर करता है

वेरिएबल-लेंथ एपिसोड्स और XLA रीकंपाइलेशन

JAX का XLA कंपाइलर फिक्स्ड टेंसर शेप्स (fixed tensor shapes) के लिए ऑप्टिमाइज़ करता है। हालाँकि, एजेंटिक टास्क अलग-अलग लंबाई के सीक्वेंस उत्पन्न करते हैं, जो सामान्य रूप से महंगे रीकंपाइलेशन को ट्रिगर करेंगे। Tunix छोटे सीक्वेंस को एक साथ पैक करता है और समान लंबाई के एपिसोड्स को बकेट्स में समूहित करता है, जिससे शेप्स इतने समय तक स्थिर रहते हैं कि XLA कंपाइल्ड कर्नेल का पुन: उपयोग कर सके। इसका परिणाम बिना किसी कंपाइलर ओवरहेड के स्थिर थ्रूपुट है, जो अन्यथा प्रदर्शन को खराब कर देता।

कई TPU चिप्स पर विशाल मॉडल्स को स्केल करना

70 बिलियन से अधिक पैरामीटर्स वाले एजेंटों को प्रशिक्षित करने के लिए वेट्स (weights) और डेटा को कई TPU नोड्स में फैलाने की आवश्यकता होती है। Tunix मॉडल पैरामीटर्स और एक्टिवेशन दोनों को शार्ड (shard) करने के लिए JAX के ShardMap प्रिमिटिव का उपयोग करता है, जिससे लर्नर को उच्च गति से डेटा फीड करते हुए पूरे मॉडल को मेमोरी में रखने की अनुमति मिलती है। यह शार्डिंग रणनीति उन मॉडल्स को प्रशिक्षित करना संभव बनाती है जो पहले एक सिंगल TPU पॉड के लिए पहुंच से बाहर थे।

डिकपल्ड पाइपलाइनों से स्टेल ग्रेडिएंट्स

जब एक्टर्स लर्नर से आगे चलते हैं, तो उनके द्वारा सप्लाई किया गया डेटा वर्तमान पॉलिसी के सापेक्ष "स्टेल" (stale) हो सकता है। Tunix इस ड्रिफ्ट को दो तंत्रों के साथ कम करता है: इम्पोर्टेंस-सैंपलिंग (importance-sampling) पुराने सैंपल्स को उनकी प्रासंगिकता दर्शाने के लिए रीवेट करती है, और एक कॉन्फ़िगर करने योग्य स्टेलनेस थ्रेशोल्ड (staleness threshold) उन ट्राजेक्टरीज़ को हटा देता है जो पूर्व-निर्धारित आयु से अधिक हो जाती हैं। साथ मिलकर वे लर्निंग को स्थिर रखते हैं, भले ही पाइपलाइन एसिंक्रोनस रूप से चल रही हो।

अपनाने वालों को किन बातों का ध्यान रखना चाहिए

  • लेटेंसी ऑडिट – डिकपलिंग का लाभ एनवायरनमेंट के रिस्पॉन्स टाइम पर निर्भर करता है। टीमों को एंड-टू-एंड लेटेंसी को मापना चाहिए और यह सुनिश्चित करना चाहिए कि बफर को अच्छी तरह से भरा रखने के लिए एक्टर पूल्स का आकार सही हो।
  • वर्कर पूल डिज़ाइन – सस्ते CPUs या GPUs कई एक्टर्स को होस्ट कर सकते हैं, लेकिन उन्हें ओवरसब्सक्राइब करने से नेटवर्क या स्टोरेज पर प्रतिस्पर्धा (contention) हो सकती है। एक संतुलित पूल जो बफर की इनजेस्ट रेट से मेल खाता हो, आवश्यक है।
  • बफर रोबस्टनेस – सेंट्रल स्टोर को नया बॉटलनेक बने बिना उच्च राइट और रीड रेट को संभालना चाहिए। कम टेल लेटेंसी (tail latency) और पर्याप्त बैंडविड्थ वाले स्टोरेज सिस्टम का चयन करना आर्किटेक्चर का एक अनिवार्य हिस्सा है।

संभावित कमियां

विभाजित आर्किटेक्चर अधिक गतिशील हिस्से (moving parts) पेश करता है: अलग-अलग हार्डवेयर फ्लीट, एक पर्सिस्टेंट बफर, और स्टेलनेस लिमिट लागू करने के लिए समन्वय तर्क (coordination logic)।

निष्कर्ष

Tunix दिखाता है कि एजेंटिक RL में प्रमुख लागत मॉडल की नहीं बल्कि सिंक्रोनस इंटरैक्शन लूप के कारण होने वाले खाली समय (idle time) की है। रोलआउट कार्य को सस्ते हार्डवेयर पर ऑफलोड करके और हाई-थ्रूपुट बफर से लगातार सीखने वाले TPU पॉड को डेटा फीड करके, Google ने 10% से कम उपयोग की समस्या को लगभग पूर्ण-क्षमता वाले वर्कफ़्लो में बदल दिया है।