Google और Meta ने TPU पर PyTorch को नेटिव रूप से चलाने के लिए TorchTPU जारी किया
टीम के अनुसार, 2025 के मध्य में Google में Meta की PyTorch टीम के साथ शुरू किया गया यह ओपन सोर्स बैकएंड पहले से ही Google Cloud के कुछ सबसे बड़े ग्राहकों द्वारा प्रोडक्शन में इस्तेमाल किया जा रहा है।
AI चित्रणGoogle और Meta ने TorchTPU पेश किया है, जो एक ओपन सोर्स बैकएंड है और PyTorch को Google के TPU चिप्स पर नेटिव रूप से चलने देता है। प्रोजेक्ट की टीम ने बताया कि यह 2025 के मध्य में Google में Meta की PyTorch टीम के सहयोग से शुरू हुआ और अभी सक्रिय विकास में है।
टीम के अनुसार, TorchTPU का इस्तेमाल पहले से ही Google Cloud के कुछ सबसे बड़े ग्राहक इन्फ़रेंस और ट्रेनिंग वर्कलोड के लिए प्रोडक्शन में कर रहे हैं। टीम ने कहा कि TPU ऐतिहासिक रूप से TensorFlow और JAX इकोसिस्टम से जुड़े रहे हैं।
टीम ने TorchTPU को पिछले प्रयासों से एक अलग राह बताया। Google ने 2020 में PyTorch/XLA जारी किया था, जो Lazy Tensors पर आधारित TPU के लिए PyTorch का पहला कम्पाइल्ड संस्करण था। टीम के अनुसार, उस इंटीग्रेशन में सोर्स कोड में दखल देने वाले बदलाव करने पड़ते थे, कम्पाइल के समय भारी लेटेंसी होती थी, PyTorch के eager एक्ज़ीक्यूशन मोड का समर्थन नहीं था और नई डिस्ट्रीब्यूटेड कंप्यूटिंग सुविधाओं में वह पीछे था। इसके विपरीत, TorchTPU eager एक्ज़ीक्यूशन को सर्वोच्च प्राथमिकता देता है और उसके रनटाइम में कोई JAX निर्भरता नहीं है।
घोषित लक्ष्य यह है कि किसी वर्कलोड को TPU पर ले जाने के लिए कोड में सिर्फ़ टारगेट डिवाइस को GPU से बदलकर TPU करना ही काफ़ी होना चाहिए। टीम ने आगाह किया कि केवल इस बदलाव से TPU का अधिकतम प्रदर्शन नहीं मिलेगा, क्योंकि किसी मॉडल के शेप, शार्डिंग और ऑपरेशन GPU के लिए ट्यून किए गए हो सकते हैं, लेकिन कहा कि मॉडल फिर भी बिना किसी अतिरिक्त सेटअप के चलना चाहिए।
टीम के अनुसार, अब तक इस प्रोजेक्ट ने ये किया है:
- eager और कम्पाइल्ड मोड का समर्थन जोड़ा, जिसमें DeferAndFuse नाम का एक नया eager मोड शामिल है;
- लगभग 450 ATen ऑपरेशनों का समर्थन जोड़ा;
- JAX/Pallas और Helion DSL के ज़रिए लिखे गए कस्टम कर्नल सक्षम किए;
- सबसे ज़्यादा डाउनलोड किए गए 5,000 Hugging Face मॉडल पर बिना अतिरिक्त सेटअप के एक्ज़ीक्यूशन को प्रमाणित किया;
- vLLM, SGLang, TorchTitan, Ray और TorchRec के साथ इंटीग्रेशन को प्रमाणित किया।
टीम ने कुछ सीमाएँ भी बताईं। एक TPU डिवाइस को केवल एक यूज़र-स्पेस प्रोसेस ही एक्सेस कर सकता है, इसलिए रीइन्फ़ोर्समेंट लर्निंग कोलोकेशन बिना अतिरिक्त सेटअप के समर्थित नहीं हो सकता, और टीम ने कहा कि वह इसके समाधान तलाश रही है। वह सीमित डायनामिज़्म का समर्थन भी जोड़ रही है, क्योंकि अंतर्निहित कम्पाइलर और रनटाइम मानते हैं कि टेंसर शेप कम्पाइल के समय ज्ञात होते हैं।
प्रिसीज़न के मामले में, TorchTPU डिफ़ॉल्ट रूप से एक मध्यम मोड का उपयोग करता है, जो bfloat16 में गुणा करता है और FP32 में संचय करता है। उच्च मोड लगभग 3 गुना या लगभग 6 गुना कंप्यूट लागत पर अधिक FP32 प्रिसीज़न वापस लाते हैं, और टीम ज़्यादातर डीप लर्निंग वर्कलोड के लिए डिफ़ॉल्ट की सिफ़ारिश करती है।