Google च्या Tunix प्रणालीमुळे मोठ्या प्रमाणावरील एजेंटिक रिइन्फोर्समेंट लर्निंग (RL) मध्ये TPUs चा कार्यक्षमतेने वापर करण्यातील अडथळा दूर झाला आहे. इंटरअॅक्शन डेटा तयार करण्याचे काम आणि पॉलिसी अपडेट करण्याचे काम वेगळे करून, Tunix TPU वापर (utilization) सिंगल-डिजिट टक्केवारीवरून जवळजवळ पूर्ण क्षमतेपर्यंत नेतो, ज्यामुळे संगणकीय संसाधनांचा (compute) अपव्यय मोठ्या प्रमाणात कमी होतो.

एजेंटिक RL मधील अडथळा

एजेंटिक RL हे अधिक परिचित असलेल्या "नेक्स्ट-टोकन" लँग्वेज-मॉडेल ट्रेनिंगपेक्षा वेगळे आहे. एका एजंटला API कॉल्स पाठवावे लागतात, कोड कार्यान्वित करावा लागतो किंवा सिम्युलेटेड एन्व्हायरनमेंटमध्ये पावले उचलावी लागतात आणि नंतर निकालावर प्रतिक्रिया द्यावी लागते. त्यामुळे ट्रेनिंग लूप सिंक्रोनस (synchronous) असतो: मॉडेल एक कृती (action) करते, एन्व्हायरनमेंट चालते, निकाल परत येतो आणि त्यानंतरच मॉडेलला ग्रेडियंट अपडेट मिळते. जेव्हा एक सिंगल एन्व्हायरनमेंट स्टेप काही सेकंद घेते, तेव्हा महागडे TPU हार्डवेअर रिकामे बसते आणि रिपोर्ट केलेला वापर १०% पेक्षा खाली जाऊ शकतो. ही अकार्यक्षमता थेट उच्च क्लाउड बिले आणि संथ संशोधन चक्रांमध्ये रूपांतरित होते.

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

Tunix दोन टप्पे—ट्रॅजेक्टरी जनरेशन आणि पॉलिसी ऑप्टिमायझेशन—वेगवेगळ्या हार्डवेअर पूलवर नेऊन या समस्येवर मात करते.

  • असिंक्रोनस ॲक्टर्स (Asynchronous actors) स्वस्त CPUs किंवा GPUs वर चालतात. प्रत्येक ॲक्टर सतत त्याच्या नियुक्त एन्व्हायरनमेंटसोबत संवाद साधतो, कृती आणि निरीक्षणे रेकॉर्ड करतो आणि resulting ट्रॅजेक्टरीज एका शेअर केलेल्या स्टोअरमध्ये स्ट्रीम करतो.
  • कंटीन्यूअस लर्नर्स (Continuous learners) समर्पित TPU Pods मध्ये असतात. लर्नर सेंट्रल बफरमधून बॅचेस घेतो आणि कोणत्याही सिंगल ॲक्टरने रोलआउट पूर्ण करण्याची वाट न पाहता ग्रेडियंट अपडेट्स करतो.
  • हाय-थ्रूपुट बफर (High-throughput buffer) मध्ये असतो, जो ट्रॅजेक्टरीजसाठी स्टेजिंग एरिया म्हणून काम करतो. लर्नर बफर डेटा पुरवतो तितक्या वेगाने वाचू शकत असल्यामुळे, TPU कधीही थांबत नाही.

याचा निव्वळ परिणाम म्हणजे एक असा ट्रेनिंग पाइपलाइन जिथे TPUs जवळजवळ सर्व वेळ व्यस्त राहतात, ज्यामुळे वापर (utilization) १००% च्या जवळ पोहोचतो.

तांत्रिक आव्हाने आणि Tunix त्यावर कशी मात करते

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

JAX चा XLA कंपायलर फिक्स्ड टेन्सर शेप्ससाठी (fixed tensor shapes) ऑप्टिमाइझ करतो. तथापि, एजेंटिक टास्क वेगवेगळ्या लांबीच्या सिक्वेन्स तयार करतात, ज्यामुळे सामान्यतः महागडे रीकंपाइलेशन ट्रिगर होऊ शकते. Tunix लहान सिक्वेन्स एकत्र पॅक करते आणि समान लांबीच्या एपिसोड्सना बकेट्समध्ये गटबद्ध करते, ज्यामुळे XLA ला कंपाईल केलेले कर्नल्स पुन्हा वापरण्यासाठी पुरेसा वेळ मिळतो आणि शेप्स स्थिर राहतात. याचा परिणाम म्हणजे कंपायलर ओव्हरहेडशिवाय स्थिर थ्रूपुट मिळतो, जो अन्यथा कामगिरीवर परिणाम करू शकला असता.

अनेक TPU चिप्सवर प्रचंड मॉडेल्स स्केल करणे

७० अब्ज पेक्षा जास्त पॅरामीटर्स असलेली एजंट्स ट्रेनिंग करण्यासाठी वेट्स (weights) आणि डेटा अनेक TPU नोड्सवर पसरवणे आवश्यक असते. Tunix मॉडेल पॅरामीटर्स आणि ॲक्टिव्हेशन्स दोन्हीसाठी JAX च्या ShardMap प्रिमिटिव्हचा वापर करते, ज्यामुळे लर्नरला संपूर्ण मॉडेल मेमरीमध्ये ठेवता येते आणि तरीही त्याला उच्च वेगाने डेटा फीड करता येतो. ही शार्डिंग स्ट्रॅटेजी अशा मॉडेल्सना प्रशिक्षित करणे शक्य करते जे पूर्वी सिंगल TPU पॉडसाठी अप्राप्य होते.

डीकपल केलेल्या पाइपलाइन्समुळे स्टेल ग्रेडियंट्स

जेव्हा ॲक्टर्स लर्नरच्या पुढे जातात, तेव्हा त्यांनी पुरवलेला डेटा सध्याच्या पॉलिसीच्या तुलनेत "स्टेल" (stale) होऊ शकतो. Tunix दोन यंत्रणांद्वारे हा फरक कमी करते: इम्पोर्टन्स-सॅम्पलिंग (importance-sampling) जुन्या सॅम्पल्सना त्यांच्या सुसंगततेनुसार पुन्हा वेट देते, आणि एक कॉन्फिगर करण्यायोग्य स्टेलनेस थ्रेशोल्ड (staleness threshold) ठरवून दिलेल्या वयापेक्षा जास्त असलेल्या ट्रॅजेक्टरीज काढून टाकते. या दोन्हीमुळे पाइपलाइन असिंक्रोनस पद्धतीने चालत असतानाही लर्निंग स्थिर राहते.

वापरकर्त्यांनी काय लक्ष ठेवणे आवश्यक आहे

  • लॅटन्सी ऑडिट (Latency audit) – डीकपलिंगचा फायदा एन्व्हायरनमेंटच्या रिस्पॉन्स टाइमवर अवलंबून असतो. टीम्सनी एंड-टू-एंड लॅटन्सी मोजली पाहिजे आणि बफर व्यवस्थित भरलेला राहण्यासाठी ॲक्टर पूल योग्य आकाराचे आहेत याची खात्री केली पाहिजे.
  • वर्कर पूल डिझाइन (Worker pool design) – स्वस्त CPUs किंवा GPUs अनेक ॲक्टर्स होस्ट करू शकतात, परंतु त्यांना ओव्हरसब्सक्राइब केल्यामुळे नेटवर्क किंवा स्टोरेजवर ताण येऊ शकतो. बफरच्या इनगेस्ट रेटशी जुळणारा संतुलित पूल आवश्यक आहे.
  • बफरची मजबूती (Buffer robustness) – सेंट्रल स्टोअरला नवीन अडथळा न बनता उच्च राइट आणि रीड रेट हाताळणे आवश्यक आहे. कमी टेल लॅटन्सी (tail latency) आणि पुरेशी बँडविड्थ असलेले स्टोरेज सिस्टम निवडणे हे आर्किटेक्चरचा एक अनिवार्य भाग आहे.

संभाव्य तोटे

विभाजित आर्किटेक्चरमुळे अधिक घटक (moving parts) वाढतात: स्वतंत्र हार्डवेअर फ्लीट्स, एक पर्सिस्टंट बफर, आणि स्टेलनेस मर्यादा लागू करण्यासाठी कोऑर्डिनेशन लॉजिक.

निष्कर्ष

Tunix हे दर्शवते की एजेंटिक RL मधील मुख्य खर्च मॉडेलचा नसून सिंक्रोनस इंटरअॅक्शन लूपमुळे निर्माण होणाऱ्या रिकाम्या वेळेचा (idle time) आहे. रोलआउटचे काम स्वस्त हार्डवेअरवर सोपवून आणि हाय-थ्रूपुट बफरमधून सतत शिकणाऱ्या TPU पॉडला डेटा देऊन, Google ने १०% पेक्षा कमी वापराच्या समस्येचे रूपांतर जवळजवळ पूर्ण क्षमतेच्या वर्कफ्लोमध्ये केले आहे.