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 పరికర కోడ్‌ను ఉపయోగించేది. స్థానిక బ్యాకెండ్ PyTorchకు చెందిన PrivateUse1 ఎక్స్‌టెన్షన్‌ను ఉపయోగించి, xm.xla_device() స్థానంలో neuron పరికరంపై ఉంచే విధానాన్ని తీసుకొస్తుందని AWS చెబుతోంది. ప్యాకేజీ పేరు పరిచయమైనదే అయినప్పటికీ, ఇప్పటికే 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 కోసం అందించారు; కలెక్టివ్, పాయింట్-టు-పాయింట్ ఆపరేషన్లు డాక్యుమెంటేషన్‌లో ఉన్నాయి. మద్దతు ఉన్న APIలలో DistributedDataParallel, FSDP2, DTensor, DeviceMeshతో పాటు పైప్‌లైన్ సమాంతరీకరణను AWS పేర్కొంటోంది. పంపిణీ శిక్షణ ప్రాజెక్టులో పరిశీలించడానికి ఇవి ఉపయోగకరమైన మార్గాలు. అయితే ఇప్పటికే ఉన్న బహుళ నోడ్‌ల జాబ్ ఎలాంటి మార్పులు లేకుండా నడుస్తుందని, లేదా నిర్దిష్ట త్రూపుట్ లక్ష్యాన్ని అందుకుంటుందని విడుదల గమనికలు నిరూపించవు.

సరిపోలే వాతావరణంతో మూల్యాంకనం మొదలవుతుంది

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 వివరణ. మార్పిడి, పనితీరు, ఆర్థిక ప్రయోజనాల వాదనలను ప్రతి వర్క్‌లోడ్‌కు ప్రత్యేకంగా ధృవీకరించాలి.