AWS च्या Neuron 2.33 च्या रिलीज नोट्स मध्ये TorchNeuron च्या सार्वजनिक बीटाची तारीख 19 ऑक्टोबर दिली आहे. त्याच 19 ऑक्टोबरची तारीख असलेली PyTorch ची घोषणा सांगते की सार्वजनिक बीटा आणि मुक्त-स्रोत आवृत्ती 20 ऑक्टोबरला उपलब्ध होते. torch-neuronx पॅकेज PyTorch मध्ये Trainium2 आणि Trainium3 साठी 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 दिले आहेत; PyTorch 2.13.0 साठी torch-neuronx 2.13.3.0 किंवा PyTorch 2.14.0 साठी torch-neuronx 2.14.3.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 ची उद्दिष्टे आणि सुरुवातीचे वापरकर्ते यांबाबतचे वर्णन. स्थलांतर, कामगिरी आणि आर्थिक लाभाच्या दाव्यांची वर्कलोडनुसार पडताळणी आवश्यक आहे.