A TorchNeuron egy új PyTorch-bővítmény, amely natív backendként teszi elérhetővé az AWS Trainiumot (Neuron) PyTorch alatt. A megoldás már támogatja az eager módot, a torch.compile használatát és a PyTorch natív elosztott primitívjeit (például DTensor, FSDP2), így a Hugging Face könyvtárai — például Transformers és Diffusers — közvetlenül képesek futni Trainiumon Optimum Neuron közbeiktatása nélkül.
Elérhetőség
A TorchNeuron nyilvános bétája 2026. október 20-tól lesz elérhető. A telepítés a következő paranccsal lesz lehetséges: pip install torch-neuronx.
Mi változik az eddigi megközelítéshez képest?
Eddig az AWS Trainium és Inferentia támogatását a Hugging Face Optimum Neuron könyvtára biztosította, amely a torch-neuronx csomagra épített és PyTorch XLA-stílusú, „compiler-szerű” (lusta tenzorok, statikus alakok, gyakori újrafordítás grafikon törés esetén) működést tett szükségessé. Optimum Neuron ezt elrejtette a felhasználók elől, de extra, Neuron-specifikus kódot tartott fenn.
TorchNeuron ehelyett közvetlenül integrálódik a PyTorch-ba és a neuron eszközt natív device-ként szolgáltatja: eager végrehajtással, PyTorch elosztott primitívekkel (DTensor, FSDP2) és torch.compile kompatibilitással. Ennek következtében a Hugging Face könyvtárai számára a Trainium mostantól olyan eszköz, mint bármelyik másik gyorsító — Transformers, Diffusers, Accelerate, PEFT és TRL ugyanazokkal az osztályokkal és script-ekkel fut rajta.
A Hugging Face javaslata mostantól az, hogy a Neuron-felhasználók közvetlenül ezeket a könyvtárakat használják Optimum Neuron helyett.
Transformers — inferencia és szöveggenerálás
Széles modellkört támogat a Transformers könyvtár natívan TorchNeuron alatt. A cikk szerint a TorchNeuronral az 100 legtöbbet használt szöveggeneráló modellarchitektúra 89%-a támogatott alapból, míg Optimum Neuron v0.4.5 esetén ez az arány 46% volt. A használat egyszerű: device_map=neuron átadásával a modell a Neuron eszközökre kerül.
A cikk példakódja bemutatja a tensor parallelizmust és a statikus cache használatát, amelyet a Neuron esetén ajánlanak, hogy elkerüljék a gyakori újrafordításokat.
Training — finomhangolás Trainiumon
Mivel a Neuron natív PyTorch device, a Trainiumon végzett tanítás ugyanazokat a Hugging Face eszközöket használja, mint más gyorsítók: Transformers, Accelerate, PEFT és TRL. Nem szükséges Neuron-specifikus trainer vagy model-fork: egy standard TRL SFTTrainer-ral készült finomhangoló script Trainiumon is lefut néhány Neuron-specifikus beállítással.
Fontos részletek a Trainiumon történő tréningnél:
- Statikus alakok: a script letiltja a torch.compile automatikus dinamikus alakjaihoz kapcsolódó viselkedést, és csomagolást (wrapped packing) használ akár 2048 tokenezésig. TRL ≥ 1.16 mellett a Neuron-beli batch-ek 128 többszörösére vannak kiegészítve, így korlátozva a különböző input-alkatokat és újrafordításokat.
- torch.compile: --torch_compile true paraméter elegendő; a torch.compile felismeri a Neuron eszközt és automatikusan a megfelelő backend-et választja.
- Async weight loading: a Transformers beépített aszinkron súlybetöltése nem támogatott Neuronon, ezért jelenleg HF_DEACTIVATE_ASYNC_LOAD=1 beállítással kell kikapcsolni.
A cikk bemutat egy futtatási példát, ahol TP_SIZE=4 és FSDP_SIZE=8 beállításokkal 32 folyamatot indítanak (4 x 8 = 32), ez megfelel 32 NeuronCores-nak. A szerzők saját méréseihez 8 Trainium2 chipet (32 NeuronCores) használtak, ami egy trn2.48xlarge példány felének felel meg.
A tanítás lefut, a trainer.save_model egy standard Transformers checkpointot ír, amely nem kötődik Neuronhoz, tehát feltölthető a Hugging Face Hubra és bármilyen más hardveren betölthető.
Diffusers — képgenerálás
A Diffusers könyvtár több pipeline-ja kompatibilis TorchNeuronnal mind eager, mind fordított (compiled) módban. Konkrét modellek, mint FLUX.1-dev, FLUX.2, Qwen-Image, Qwen-Image-2.1 és H3 kompatibilisek tensor-parallelismussal is, hogy kihasználják a Neuron eszközöket.
A cikkben szerepel egy példakód, amely TP=8 konfigurációban mutatja be a Qwen-Image-2.1 pipeline használatát: a DeviceMesh inicializálása, a transformer párhuzamos konfigurálása, a model.compile_repeated_blocks meghívása Neuron backenddel, majd a kép generálása.
Futtatási példa: NEURON_LOGICAL_NC_CONFIG=2 torchrun --nproc_per_node=8 run.py. A cikk jelezte, hogy egyre több pipeline lesz TP-kompatibilis, és ha egy adott pipeline nem támogatott, a fejlesztők követhetik egy útmutatót a támogatás hozzáadásához.
Kernels és saját Neuron-műveletek
A Neuron csapat által készített kernel-ek egyre nagyobb számban jelennek meg a Hubon, hogy felgyorsítsák a modelleket. A Neuron Kernel Interface (NKI) lehetővé teszi egyedi műveletek implementálását és megosztását a Hubon. A cikk említ egy példakernelt, amelykel vision transformer-ek akár 2.4x gyorsulást érhetnek el Trainiumon.
Köszönetnyilvánítás
A cikk külön köszönetet mondott az AWS csapatának a TorchNeuron fejlesztéséért és a Hugging Face integráció javításához nyújtott gyors visszajelzésekért.
Összegzés
TorchNeuron lényegében egyszerűsíti a Trainiumon való munkát azáltal, hogy eltünteti a korábbi „compiler-szerű” réteget és natív PyTorch élményt ad: eager végrehajtás, torch.compile támogatás és a megszokott Hugging Face könyvtárak futtatása. A 2026. október 20-i public beta indulással a kutatók és fejlesztők könnyebben telepíthetik és megoszthatják modelljeiket Trainiumon a Hubon keresztül.
