TorchNeuronin beeta avasi PyTorchin Trainiumille

AWS avasi TorchNeuronin julkisen beetan 20. lokakuuta 2026 Amazon EC2:n Trainium2- ja Trainium3-instansseille. Avoimen lähdekoodin natiivi PyTorch-taustajärjestelmä antaa kehittäjälle tutun tavan ajaa mallia Trainiumilla: malli ja sen syötteet siirretään neuron-laitteelle esimerkiksi .to('neuron')-kutsulla.

Julkisessa beetassa malli voi toimia PyTorchin eager-tilassa tai käyttää torch.compile-käännöstä. Laitemäärityksen vaihto madaltaa kokeilun kynnystä, mutta se ei yksin ratkaise yhteensopivuutta. Mallin käyttämät rajapinnat, syötteiden muodot ja suorittimelle mahdollisesti siirtyvät operaatiot määräävät, kuinka paljon siirto lopulta vaatii työtä.

Versioiden ja asennusympäristön on sovittava yhteen

AWS:n TorchNeuronin julkaisutiedot yhdistävät PyTorch 2.13.0:n torch-neuronx-pakettiperheeseen 2.13.3.0.* ja PyTorch 2.14.0:n pakettiperheeseen 2.14.3.0.*. Molemmille on ilmoitettu Python-versiot 3.12, 3.13 ja 3.14. Taustajärjestelmä rekisteröi neuron-laitteen PyTorchin PrivateUse1-laajennuksella, joten myös syötetensorien on oltava samalla laitteella kuin malli.

Yhteensopivuus riippuu lisäksi asennustavasta. Julkaisu tukee Linuxin x86_64-ympäristöä Ubuntu 22.04-, 24.04- ja 26.04-versioilla, mutta tämän julkaisun erilliset wheel-paketit ovat saatavilla vain Ubuntu 26.04:lle. Ubuntu 24.04:lle on tarjolla valmis DLAMI-ympäristö; vanhempien Ubuntu-versioiden manuaalinen asennus edellyttää lähdekoodista rakentamista.

Uusi natiivi taustajärjestelmä muuttaa myös aiemmasta Trainium-toteutuksesta siirtyvän sovelluksen laitekutsuja. XLA-pohjaiseen torch-neuronx-toteutukseen kirjoitetut torch_xla- ja xm.xla_device()-kutsut eivät ole uuden neuron-laitteen käyttötapa. Siksi valmiin mallin siirrossa on erotettava tavallinen laitemäärityksen vaihto vanhan Trainium-koodin muuttamisesta.

Tuettujen operaatioiden tietotyyppeihin kuuluvat muun muassa bfloat16, float16 ja float32, ja torch.autocast tukee sekoitettua tarkkuutta neuron-laitteella. float64 käsitellään kuitenkin float32-tarkkuudella ja siitä annetaan varoitus. Jos mallin laskenta edellyttää kaksinkertaista tarkkuutta, pelkkä onnistunut ajo ei riitä osoittamaan, että numeerinen käyttäytyminen säilyy ennallaan.

Eager-ajo ja torch.compile palvelevat eri vaiheita

Eager-tilassa operaatiot suoritetaan Python-kutsujen edetessä. Tavallinen ohjausvirta, tulosteet ja keskeytyspisteet ovat siten käytettävissä kehitystyössä. Ensimmäinen operaatiokutsu voi silti sisältää käännöksen; vastaava käännetty tulos voidaan käyttää myöhemmin välimuistista. Ensimmäisen ajon viive ei siksi kuvaa yksin toistuvan ajon nopeutta.

torch.compile voi kääntää koko mallin tai valittuja alueita neuron-taustajärjestelmälle, myös tuettuja takaisinkulun laskentakaavioita. Käännetyn alueen ulkopuolinen koodi jatkaa eager-tilassa. Beetassa käännös käyttää staattisia syötemuotoja ja vaatii asetuksen dynamic=False; uusi syötemuoto voi käynnistää uuden käännöksen.

Vaihtelevan pituisessa tekstissä tai muuttuvilla eräko'oilla ero on käytännöllinen. Yhdellä syötemuodolla onnistunut ajo ei kerro, kuinka usein todellinen kuorma tuottaa uusia käännöksiä. Syötteiden täyttäminen tai ryhmittely vakiomuotoihin voi vähentää vaihtelua. Käännetyt artefaktit tallentuvat levylle, ja myös eri prosessit voivat käyttää sopivaa välimuistimerkintää.

Tuettu hajautettu ajo ulottuu tavallisiin PyTorch-rajapintoihin, kuten DistributedDataParalleliin, fully_shard-toimintoon, DTensoriin ja DeviceMeshiin. Niillä voidaan jakaa koulutusta laitteiden kesken ilman erillistä mallin hajautusrajapintaa. Neuron Kernel Interface puolestaan mahdollistaa omien operaattorien toteuttamisen, jos mallin tärkeä laskentavaihe vaatii erillisen laitetoteutuksen.

API-aukot ja CPU-varatoiminnot rajaavat beetaa

Julkinen beeta ei tue torch.sparse-, torch.func-, torch.nested- eikä torch.distributions-alueita. Näihin tukeutuva malli voi siis tarvita muutoksia, vaikka sen tavalliset tensoritoiminnot toimisivat neuron-laitteella. Myös torch.fft-tuki on osittainen: esimerkiksi torch.fft.rfft toimii laitteella, kun taas muut FFT-muunnokset suoritetaan suorittimella.

Eager-tilassa tuen ulkopuolelle jäävä operaatio voi siirtyä CPU:lle. Malliajon onnistuminen ei siten osoita, että koko laskenta tapahtui Trainiumilla. torch_neuronx.get_fallback_ops() näyttää kirjatut varatoiminnot, ja torch.profiler voi erottaa laitteen työn suorittimella tehdystä työstä. Tämä on tärkeää sekä viiveen arvioinnille että operaattorien muutostarpeelle.

Kvantisointiin liittyvät rajat koskevat mallin toteutusta suoraan: julkaisu ei sisällä TorchAO-integraatiota eikä yhdistetyn MXFP8-tietotyypin tukea. Jos nykyinen malli käyttää näitä ominaisuuksia, sen kvantisoitua polkua ei voi olettaa siirtyväksi sellaisenaan. Myös laitteen tukema mutta eri tarkkuudella suoritettu laskenta vaatii tulosten erillistä vertailua.

Yhteensopivuus jakautuu näin kolmeen eri kysymykseen: käynnistyykö malli, missä sen operaatiot todella suoritetaan ja vastaavatko tulokset aiempaa ympäristöä. Natiivi PyTorch-laite ratkaisee ensimmäisen askeleen monelle mallille, mutta API-aukot ja CPU-varatoiminnot voivat muuttaa kahden muun vastauksen. Beetasta ei ole esitetty yleispätevää suorituskykytulosta kaikille malleille.

Hugging Facen kirjastot laajentavat käyttöä

Hugging Facen TorchNeuron-julkaisun mukaan Transformers-, Diffusers-, Accelerate-, PEFT- ja TRL-kirjastoja voi käyttää suoraan neuron-laitteen kanssa aiemman Optimum Neuron -välikerroksen sijasta. Sen vertailussa TorchNeuron tukee suoraan 89:ää sadasta käytetyimmästä tekstinluontimallin arkkitehtuurista, kun Optimum Neuron v0.4.5:n osuus on 46 prosenttia. Luvut mittaavat tämän arkkitehtuurijoukon yhteensopivuutta, eivät mallien ajonopeutta.

Tekstinluonnin esimerkissä käytetään staattista välimuistia toistuvien käännösten välttämiseksi. Koulutusesimerkissä taas automaattiset dynaamiset muodot poistetaan käytöstä, ja sekvenssit täytetään ennakoitaviin pituuksiin. Kirjastot voivat siis säilyttää tutut luokat ja koulutusskriptit, vaikka Trainium-ajolle tarvitaan edelleen syötemuotoihin liittyviä asetuksia.

Laajentunut kirjastotuki tekee ensimmäisestä kokeilusta pienemmän koodimuutoksen, mutta ei poista mallikohtaista validointia. Erityisesti operaattorien CPU-siirtymät, käännösten määrä ja numeeriset tulokset ratkaisevat, vastaako toimiva TorchNeuron-ajo myös kyseisen mallin käyttötarvetta.

Lue myös: