AWS के Neuron 2.33 रिलीज़ नोट्स में TorchNeuron के सार्वजनिक बीटा की तारीख 19 अक्टूबर दी गई है। 19 अक्टूबर की ही तारीख वाली PyTorch की घोषणा कहती है कि सार्वजनिक बीटा और ओपन-सोर्स रिलीज़ 20 अक्टूबर को उपलब्ध होगी। torch-neuronx पैकेज Trainium2 और Trainium3 के लिए PyTorch में एक neuron डिवाइस पंजीकृत करता है, जिससे समर्थित मॉडल और टेंसर ऑपरेशन परिचित PyTorch API का उपयोग कर सकते हैं। इन चिप पर विचार कर रहे इंजीनियरों के लिए उपयोगी बदलाव यह है कि मौजूदा PyTorch प्रोग्राम से डिवाइस पर परीक्षण तक का रास्ता छोटा हो जाता है। प्रोग्राम सही और कुशलता से चलेगा या नहीं, यह अब भी उसके ऑपरेशन, आकार और सॉफ़्टवेयर वातावरण पर निर्भर करता है।

Trainium के लिए पहले का PyTorch रास्ता XLA-आधारित torch-neuronx लाइब्रेरी और torch_xla डिवाइस कोड का उपयोग करता था। AWS के अनुसार, नेटिव बैकएंड PyTorch का PrivateUse1 एक्सटेंशन उपयोग करता है और xm.xla_device() की जगह neuron डिवाइस पर रखने की प्रक्रिया लाता है। पैकेज का नाम परिचित होने के बावजूद, पहले से Neuron उपयोग करने वाले प्रोजेक्ट के लिए यह इंटरफ़ेस में महत्त्वपूर्ण बदलाव है। यह रिलीज़ Trn2 और Trn3 इंस्टेंस पर ट्रेनिंग और इन्फ़रेंस के लिए है; समर्थन तालिका में Trn1 शामिल नहीं है।

PyTorch प्रोग्राम में क्या बदलता है

इन AWS रिलीज़ नोट्स में मॉडल को .to("neuron") से स्थानांतरित करना और इनपुट टेंसर को device="neuron" पर बनाना दिखाया गया है। समर्थित torch.nn, autograd और ऑप्टिमाइज़र ऑपरेशन फिर बैकएंड के ज़रिये चलते हैं। PyTorch के डिवाइस-निरपेक्ष torch.accelerator API के लिए लिखा गया कोड किसी खास डिवाइस के लिए दोबारा लिखे बिना Neuron को खोज सकता है, जैसा कि TorchNeuron के अवलोकन में बताया गया है।

चलाने के दो रास्ते हैं। ईगर मोड Python के बुलाते ही PyTorch ऑपरेशन चलाता है और सामान्य कंट्रोल फ़्लो, प्रिंट आउटपुट तथा ब्रेकपॉइंट बनाए रखता है। हर ऑपरेशन पहली बार उपयोग पर कंपाइल होता है और बाद के मेल खाते कॉल कैश का उपयोग कर सकते हैं। इंजीनियर मॉडल या चुने हुए हिस्से को torch.compile(model, backend="neuron", dynamic=False) से भी कंपाइल कर सकते हैं। नेटिव बैकएंड समर्थित फ़ॉरवर्ड और बैकवर्ड ग्राफ़ कंपाइल करता है; कैप्चर किए गए हिस्से के बाहर का कोड ईगर मोड में चलता रहता है। PyTorch की घोषणा कंपाइल करने से पहले torch.compiler.set_default_backend("neuron") को एक बार सेट करने का तरीका भी बताती है, ताकि बाद के कंपाइल कॉल में बैकएंड आर्ग्युमेंट छोड़ा जा सके।

रिलीज़ में एक neuron बैकएंड torch.distributed के लिए शामिल है, जिसके कलेक्टिव और पॉइंट-टू-पॉइंट ऑपरेशन दस्तावेज़ों में दिए गए हैं। AWS समर्थित API में DistributedDataParallel, FSDP2, DTensor, DeviceMesh और पाइपलाइन समानांतरता को सूचीबद्ध करता है। वितरित ट्रेनिंग प्रोजेक्ट में इन रास्तों की जाँच उपयोगी है, लेकिन रिलीज़ नोट्स यह स्थापित नहीं करतीं कि मौजूदा मल्टी-नोड जॉब बिना बदलाव चलेगा या किसी खास थ्रूपुट लक्ष्य को हासिल करेगा।

मूल्यांकन मेल खाते वातावरण से शुरू होता है

AWS की इंस्टॉलेशन की त्वरित गाइड के लिए Trn2 या Trn3 इंस्टेंस तथा Neuron Base DLAMI या समर्थित Ubuntu AMI आवश्यक है। रिलीज़ तालिका Ubuntu 24.04 या 26.04 पर Linux x86_64 निर्दिष्ट करती है। इस रिलीज़ में Ubuntu 22.04 का समर्थन समाप्त हो चुका है। इसमें Python 3.12, 3.13 और 3.14 सूचीबद्ध हैं, तथा torch-neuronx 2.13.3.0 PyTorch 2.13.0 के लिए या torch-neuronx 2.14.3.0 PyTorch 2.14.0 के लिए है। पैकेज परिवार को PyTorch संस्करण से मेल खाना चाहिए।

त्वरित गाइड में torch-neuronx==2.14.* के साथ neuronx-cc==2.* और nki को Neuron और PyTorch पैकेज इंडेक्स से इंस्टॉल किया जाता है। फिर डिवाइस की उपलब्धता जाँची जाती है और Neuron पर एक छोटे मैट्रिक्स ऑपरेशन की तुलना CPU के नतीजे से होती है। यह AWS के उदाहरण में इंस्टॉल किए गए डिवाइस पथ की पुष्टि करता है। इससे पाठक के मॉडल की पुष्टि नहीं होती। अगला समझदारी भरा मूल्यांकन वास्तविक मॉडल के फ़ॉरवर्ड और बैकवर्ड पथ चलाना, torch_neuronx.get_fallback_ops() से दर्ज CPU फ़ॉलबैक देखना, मौजूदा वातावरण से आउटपुट की तुलना करना और कंपाइलेशन व स्थिर अवस्था के रन अलग-अलग मापना है। ये अंतिम कदम BIG CHANGE के सुझाए मूल्यांकन प्रश्न हैं, हमारे किए हुए परीक्षण नहीं।

बीटा अब भी काम को कहाँ बदल सकता है

ईगर मोड में असमर्थित ऑपरेशन CPU पर फ़ॉलबैक कर सकता है और चेतावनी दे सकता है। torch.compile के तहत, ग्राफ़ में कैप्चर किए गए ऑपरेशन को Neuron बैकएंड निचले स्तर के निष्पादन रूप में नहीं बदल पाए, तो डिफ़ॉल्ट रूप से कंपाइलेशन विफल होता है। हर ऑपरेशन के समर्थन की जानकारी के लिए AWS डेवलपर को अपनी ऑपरेटर अनुकूलता तालिका की ओर भेजता है। इसलिए ईगर रन पूरा होना अपने आप में पूरे काम के डिवाइस पर चलने या कंपाइल किए गए निष्पादन की अनुकूलता का प्रमाण नहीं है।

रिलीज़ यह भी कहती है कि torch.compile के लिए dynamic=False आवश्यक है; इनपुट का नया आकार एक और कंपाइलेशन माँग सकता है। आकार की विविधता कम करने के लिए AWS पैडिंग या बकेटिंग सुझाता है। मेल खाता आर्टिफ़ैक्ट कैश में न हो, तो पहले रन में कंपाइलेशन शामिल होता है। वर्कलोड का समय मापते हुए ये बातें महत्त्वपूर्ण हैं: कंपाइलेशन का विलंब और इनपुट आकारों का पैटर्न मूल्यांकन का परिणाम बदल सकते हैं।

अन्य स्पष्ट सीमाएँ अनुकूलता के दावे का दायरा घटाती हैं। बीटा torch.sparse, torch.func, torch.nested या torch.distributions का समर्थन नहीं करता। केवल कुछ torch.fft ट्रांसफ़ॉर्म Neuron पर चलते हैं, और float64 को float32 के रूप में चलाया जाता है, साथ में चेतावनी आती है। इस रिलीज़ में TorchAO एकीकरण और संयुक्त MXFP8 डेटा प्रकार का समर्थन नहीं है। इन इंटरफ़ेस का उपयोग करने वाली टीम को डिवाइस पर रखने को माइग्रेशन योजना मानने से पहले अपना कोड पथ जाँचना होगा।

बड़ा बदलाव

TorchNeuron सार्वजनिक बीटा में Trainium को PyTorch के नेटिव डिवाइस और कंपाइलेशन इंटरफ़ेस के पीछे लाता है। इससे समर्थित PyTorch प्रोग्राम के लिए पहला परीक्षण अधिक सीधा होता है और मौजूदा Neuron उपयोगकर्ताओं को उनके XLA-आधारित कोड से अलग रास्ता मिलता है। अपनाने का निर्णायक प्रमाण हर वर्कलोड से आएगा: क्या उसके ऑपरेटर अपेक्षित डिवाइस पर रहते हैं, क्या संख्यात्मक परिणाम स्वीकार्य बने रहते हैं, और क्या मापा गया प्रदर्शन व लागत इंस्टेंस तथा माइग्रेशन के काम को उचित ठहराते हैं। इस विश्लेषण के लिए BIG CHANGE ने दस्तावेज़ और घोषणाएँ देखीं; हमने TorchNeuron नहीं चलाया या AWS के प्रदर्शन संबंधी दावों की स्वतंत्र पुष्टि नहीं की।

स्रोत और आगे पढ़ने के लिए

  • AWS Neuron 2.33 TorchNeuron रिलीज़ नोट्स: रिलीज़ की तारीख, बीटा स्थिति, संस्करण व ऑपरेटिंग-सिस्टम तालिका, API और ज्ञात सीमाएँ।
  • AWS का TorchNeuron अवलोकन: डिवाइस एकीकरण, ईगर और कंपाइल किया गया निष्पादन, CPU फ़ॉलबैक का व्यवहार और Neuron के ज़रिये सॉफ़्टवेयर पथ।
  • AWS की इंस्टॉलेशन और सत्यापन की त्वरित गाइड: इंस्टेंस व इमेज की पूर्वापेक्षाएँ, पैकेज इंस्टॉलेशन और दस्तावेज़ों में दिया गया डिवाइस परीक्षण।
  • PyTorch ब्लॉग: TorchNeuron का परिचय: डिवाइस बदलने, कंपाइलेशन और वितरित API की PyTorch की ओर से व्याख्या। यह लेख PyTorch ब्लॉग पर AWS योगदानकर्ताओं की घोषणा है, कोई स्वतंत्र बेंचमार्क नहीं।
  • Amazon की घोषणा: अपने लक्ष्यों और शुरुआती उपयोगकर्ताओं का AWS का विवरण। माइग्रेशन, प्रदर्शन और आर्थिक लाभ के दावों को हर वर्कलोड के लिए अलग से सत्यापित करना होगा।