TorchNeuron vereinfacht HF-Modellbetrieb auf AWS Trainium

TorchNeuron vereinfacht die Nutzung von Hugging-Face-Modellen auf AWS Trainium erheblich. Ab dem 20. Oktober 2026 steht das Plugin als öffentliche Beta zur Verfügung und kann über pip install torch-neuronx installiert werden. TorchNeuron integriert Trainium als natives PyTorch-Backend und unterstützt Eager Execution, torch.compile sowie PyTorch-eigene verteilte API-Funktionen wie FSDP2, DTensor und Tensor Parallelismus. Damit entfällt der bisherige Umweg über Optimum Neuron und die auf PyTorch XLA basierende Compiler-Logik mit Lazy Tensors, statischen Formen und häufiger Neuberechnung von Graphbreaks. Für die Hugging-Face-Bibliotheken Transformers, Diffusers, Accelerate, PEFT und TRL wird Trainium dadurch wie ein gewöhnlicher Beschleuniger behandelt. Bisherige Nutzer sollen die Standardbibliotheken direkt verwenden, statt Optimum Neuron. Für Inferenzaufgaben mit Textgenerierung unterstützt TorchNeuron nach offiziellen Angaben 89 Prozent der 100 meistgenutzten Modellarchitekturen nativ. Zuvor lag der Abdeckungsgrad in Optimum Neuron 0.4.5 bei 46 Prozent. Modelle können in Transformers mit device_map=neuron auf Trainium geladen werden. Für stabile Ergebnisse empfehlen die Herausgeber statischen Caching-Code beim Generieren, da häufige Rekompilierungen vermieden werden sollen. Tensor Parallelismus lässt sich über die verteilten PyTorch-Primitives direkt nutzen, um Neuron-Instanzen vollständig auszunutzen. Auch das Training auf Trainium wird einheitlicher. Statt spezieller Trainer oder Neuron-spezifischer Modellklassen können Standard-Skripte aus Transformers, Accelerate, PEFT und TRL verwendet werden. Verteiltes Training wird über eine Distributed-Config gesteuert, in der Tensor-Parallelismus und FSDP-Größe gesetzt werden. Das Modell wird bereits beim Laden über ein TP-FSDP-Gitter geschachtelt, sodass der Trainer es nicht erneut verpackt. Ein Beispiel auf 8 Trainium2-Chips mit 32 NeuronCores zeigt, dass Standard-TRL-Skripte mit wenigen Neuron-spezifischen Einstellungen trainieren können. Die daraus entstehenden Checkpoints sind hardwareunabhängig und lassen sich ohne Anpassungen auf dem Hugging Face Hub veröffentlichen. Diffusers unterstützt ebenfalls viele Bildgenerierungspipelines auf TorchNeuron, sowohl im Eager-Modus als auch mit torch.compile. Bereits jetzt sind Pipelines wie FLUX.1-dev, FLUX.2, Qwen-Image, Qwen-Image-2.1 und H3 mit Tensor Parallelismus kompatibel, um mehrere Neuron-Geräte parallel zu nutzen. Die Integration erfolgt über Standard-PyTorch-Objekte und Neuron-spezifische Compile-Backends. Der Textencoder kann dabei weiterhin auf der CPU bleiben, während die Transformer-Komponenten auf Trainium verteilt und kompiliert werden. Weitere Pipelines sollen schrittweise folgen, und Nutzer können fehlende Unterstützung mit einer bereitgestellten Anleitung selbst ergänzen. Zusätzlich öffnet die Hugging-Face-Plattform Neuron-Kernel für Modelle. Diese können über die Neuron Kernel Interface geschrieben werden und ermöglichen eigene optimierte Operationen. Kernel lassen sich wie weitere Modellbausteine auf dem Hub teilen. Nach Angaben der Teammitglieder können einzelne Kernel Geschwindigkeitsgewinne von bis zum Faktor 2,4 auf Trainium erzielen. Die enge Zusammenarbeit mit AWS soll die Integration von TorchNeuron in die Hugging-Face-Infrastruktur dauerhaft verbessern und Trainium als gleichberechtigtes Zielsystem für Training, Feinabstimmung und Inferenz etablieren.