TorchNeuron: natív PyTorch-eszköz lett az AWS Trainium

Mi változik?

Eddig a Hugging Face modelljei az optimum-neuron könyvtáron keresztül futottak Trainiumon és Inferentián. Ez a PyTorch XLA-ra épülő torch-neuronx csomagot használta, vagyis lazy tensorokat, kizárólag statikus alakokat és minden graph breaknél újrafordítást. Az Optimum Neuron ezt a bonyolultságot rejtette el saját modellosztályokkal és trainerekkel.

A TorchNeuron a PyTorch hivatalos PrivateUse1 bővítési pontján regisztrálja a neuron eszközt. Támogatja az eager módot (kísérleti adaptive eager móddal együtt), a torch.compile-t, a natív elosztott API-kat (DTensor, FSDP2, tensor parallelism), illetve az NKI-val írt saját kerneleket. Trainium2-n és Trainium3-on fut, Apache 2.0 licenccel, a telepítés pedig ennyi: pip install torch-neuronx (PyTorch).

Mit jelent ez a Hugging Face-felhasználóknak?

A Transformers, a Diffusers, az Accelerate, a PEFT és a TRL közvetlenül fut Trainiumon, ugyanazokkal az osztályokkal és scriptekkel, mint GPU-n. A Hugging Face mostantól azt javasolja, hogy a Neuron-felhasználók ezeket a könyvtárakat használják az Optimum Neuron helyett (Hugging Face).

• Inferencia: a 100 leggyakrabban használt szöveggeneráló architektúra 89%-a azonnal működik, szemben az Optimum Neuron v0.4.5 46%-ával. Elég a device_map="neuron" paraméter.
• Tanítás: egy szokásos TRL SFTTrainer script néhány Neuron-specifikus beállítással fut. A HF 8 Trainium2 chipen (32 NeuronCore, egy trn2.48xlarge fele) tesztelte, a kimenet pedig sima Transformers checkpoint, ami bármilyen hardveren betölthető.
• Diffusers: sok pipeline fut eager és compiled módban, a FLUX.1-dev, a FLUX.2, a Qwen-Image, a Qwen-Image-2.1 és a H3 tensor parallelismmel is.
• Kernels: a Hubon NKI-kernelek jelennek meg, a HF szerint egy vision transformer kernel akár 2,4-szeres gyorsulást ad Trainiumon.

Mire kell figyelni?

A statikus alakok továbbra is számítanak. Generálásnál static cache ajánlott, tanításnál a scriptek kikapcsolják a torch.compile automatikus dinamikus alakjait, a TRL 1.16-tól pedig 128 token többszörösére paddeli a batcheket. A Transformers aszinkron súlybetöltése Neuronon egyelőre nem megy, ezt a HF_DEACTIVATE_ASYNC_LOAD=1 változóval kell kikapcsolni.

A gyakorlati haszon egyértelmű: aki Trainiumon akar tanítani vagy inferálni, annak nem kell külön kódágat fenntartania. Mivel a csomag még beta, éles terhelés előtt érdemes a saját modellen kimérni a teljesítményt.