AWS-ove beleške o izdanju Neuron 2.33 navode 19. oktobar kao datum javne beta verzije TorchNeuron-a. PyTorch objava, takođe datirana na 19. oktobar, kaže da javna beta verzija i izdanje otvorenog koda postaju dostupni 20. oktobra. Paket torch-neuronx registruje neuron uređaj u PyTorch-u za Trainium2 i Trainium3, pa podržane operacije nad modelima i tenzorima mogu da koriste poznate PyTorch API-je. Za inženjere koji razmatraju ove čipove, korisna promena je kraći put od postojećeg PyTorch programa do probe na uređaju. Da li će taj program raditi ispravno i efikasno i dalje zavisi od njegovih operacija, oblika i softverskog okruženja.

Raniji PyTorch put za Trainium koristio je biblioteku torch-neuronx zasnovanu na XLA i kod za uređaj torch_xla. AWS kaže da izvorni backend koristi PyTorch proširenje PrivateUse1 i zamenjuje xm.xla_device() smeštanjem na uređaj neuron. To je značajna promena interfejsa za projekte koji već koriste Neuron, iako je naziv paketa poznat. Izdanje važi za obuku i inferenciju na instancama Trn2 i Trn3; matrica podrške ne navodi Trn1.

Šta se menja u PyTorch programu

U AWS-ovim beleškama o izdanju prikazan je model premešten pomoću .to("neuron") i ulazni tenzor napravljen uz device="neuron". Podržane operacije torch.nn, autograd-a i optimizatora zatim se izvršavaju preko backend-a. Kod napisan za PyTorch API nezavisan od uređaja torch.accelerator može da otkrije Neuron bez prerade specifične za taj uređaj, prema dokumentu Pregled TorchNeuron-a.

Postoje dva puta izvršavanja. Eager režim izvršava PyTorch operacije kada ih Python pozove, uz očuvanje uobičajenog toka kontrole, ispisa i prekidnih tačaka. Svaka operacija se kompajlira pri prvom korišćenju, a naredni odgovarajući pozivi mogu da koriste keš. Inženjeri mogu i da kompajliraju model ili izabrani deo pomoću torch.compile(model, backend="neuron", dynamic=False). Izvorni backend kompajlira podržane grafove prolaza unapred i unazad; kod van obuhvaćenog dela nastavlja da radi u eager režimu. PyTorch objava opisuje i jednokratno postavljanje torch.compiler.set_default_backend("neuron") pre kompilacije, tako da kasniji pozivi kompilacije mogu da izostave argument backend-a.

Izdanje uključuje neuron backend za torch.distributed, sa dokumentovanim kolektivnim operacijama i operacijama od tačke do tačke. AWS među podržanim API-jima navodi DistributedDataParallel, FSDP2, DTensor i DeviceMesh, kao i cevovodni paralelizam. To su korisni putevi za razmatranje u projektu distribuirane obuke, ali beleške o izdanju ne dokazuju da će postojeći zadatak na više čvorova raditi bez izmena ili dostići određeni cilj propusne moći.

Evaluacija počinje usklađenim okruženjem

AWS-ovo kratko uputstvo za instalaciju zahteva instancu Trn2 ili Trn3 i Neuron Base DLAMI ili podržani Ubuntu AMI. Matrica izdanja navodi Linux x86_64 na Ubuntu 24.04 ili 26.04. Podrška za Ubuntu 22.04 je završena za ovo izdanje. Navedeni su Python 3.12, 3.13 i 3.14, uz torch-neuronx 2.13.3.0 za PyTorch 2.13.0 ili torch-neuronx 2.14.3.0 za PyTorch 2.14.0. Porodica paketa mora da odgovara verziji PyTorch-a.

Kratko uputstvo instalira torch-neuronx==2.14.* zajedno sa neuronx-cc==2.* i nki iz Neuron i PyTorch indeksa paketa. Zatim proverava dostupnost uređaja i poredi malu matričnu operaciju na Neuron-u sa rezultatom na CPU-u. Time se proverava put do instaliranog uređaja u AWS-ovom primeru. Ne proverava se čitaočev model. Razuman sledeći korak evaluacije je da se izvrše prolazi unapred i unazad stvarnog modela, pregledaju zabeležene operacije prebačene na CPU pomoću torch_neuronx.get_fallback_ops(), uporede izlazi sa trenutnim okruženjem i odvojeno izmere kompilacija i izvršavanje u ustaljenom stanju. Ti poslednji koraci su pitanja za evaluaciju koja predlaže BIG CHANGE, a ne testovi koje smo mi sproveli.

Gde beta verzija još može da promeni rad

U eager režimu, nepodržana operacija može da pređe na CPU i prikaže upozorenje. Kada se koristi torch.compile, operacija obuhvaćena grafom koju Neuron backend ne može da prevede u niži izvršni oblik podrazumevano dovodi do neuspešne kompilacije. AWS upućuje programere na svoju tabelu kompatibilnosti operatora za pokrivenost svake operacije. Zato uspešno završeno izvršavanje u eager režimu samo po sebi ne dokazuje da je sve izvršeno na uređaju ili da je podržana kompilacija.

U izdanju se takođe navodi da torch.compile zahteva dynamic=False; novi oblik ulaza može da zahteva novu kompilaciju. AWS predlaže dopunjavanje ili grupisanje po oblicima kako bi se smanjile njihove varijacije. Prvo izvršavanje uključuje kompilaciju ako nema odgovarajućeg artefakta u kešu. Ovi detalji su važni pri merenju radnog opterećenja: kašnjenje kompilacije i obrazac oblika ulaza mogu da promene rezultat evaluacije.

Druga izričita ograničenja sužavaju tvrdnju o kompatibilnosti. Beta verzija ne podržava torch.sparse, torch.func, torch.nested ili torch.distributions. Samo neke torch.fft transformacije rade na Neuron-u, a float64 se izvršava kao float32 uz upozorenje. Ovo izdanje nema integraciju sa TorchAO niti podršku za složeni tip podataka MXFP8. Tim koji koristi ove interfejse mora da proveri sopstveni put kroz kod pre nego što smeštanje na uređaj shvati kao plan migracije.

Velika promena

TorchNeuron u javnoj beta verziji postavlja Trainium iza PyTorch-ovih izvornih interfejsa za uređaje i kompilaciju. Time prva proba postaje neposrednija za podržane PyTorch programe, a postojeći korisnici Neuron-a dobijaju drugačiji put od svog koda zasnovanog na XLA. Presudni dokazi za usvajanje dolaziće iz svakog radnog opterećenja: da li njegovi operatori ostaju na predviđenom uređaju, da li numerički rezultati ostaju prihvatljivi i da li izmerene performanse i troškovi opravdavaju instancu i rad na migraciji. BIG CHANGE je za ovu analizu pregledao dokumentaciju i objave; nismo pokretali TorchNeuron niti nezavisno proveravali AWS-ove tvrdnje o performansama.

Izvori i dodatno čitanje