TorchNeuron: Hugging Face-Modelle laufen nativ auf AWS Trainium

Trainium wird ein normales PyTorch-Device

Bisher lief Hugging Face auf AWS Trainium und Inferentia über optimum-neuron, das auf PyTorch XLA basiert: lazy Tensors, nur statische Shapes, Neukompilierung bei jedem Graph-Break. Mit TorchNeuron, einem PyTorch-Plugin, ändert sich das. Es stellt „neuron“ als natives Device bereit, mit Eager Mode, torch.compile sowie DTensor und FSDP2. Ab dem 20. Oktober 2026 gibt es eine öffentliche Beta (pip install torch-neuronx).

Was das für Hugging Face heißt

• Transformers, Diffusers, Accelerate, PEFT und TRL laufen direkt auf Trainium, mit denselben Klassen und Skripten wie auf anderen Beschleunigern. Hugging Face empfiehlt, optimum-neuron nicht mehr zu nutzen.
• Textgenerierung: 89 % der 100 meistgenutzten Architekturen laufen ohne Anpassung, gegenüber 46 % mit Optimum Neuron v0.4.5. Es genügt device_map=neuron.
• Training: Ein normales TRL-SFTTrainer-Skript läuft mit wenigen Neuron-Einstellungen. Tensor-Parallelismus und FSDP werden über eine DistributedConfig in Transformers gesteuert, ohne eigenen Trainer.
• Hinweise: statische Shapes (Batches auf Vielfache von 128 Tokens gepolstert), torch.compile wählt das Neuron-Backend selbst, asynchrones Gewichtsladen ist noch nicht unterstützt (HF_DEACTIVATE_ASYNC_LOAD=1).

Das Beispiel lief auf 8 Trainium2-Chips (32 NeuronCores) mit TP=4 und FSDP=8.

Quelle: Hugging Face Blog