TorchNeuron macht die Ausführung von Hugging Face-Modellen auf AWS Trainium einfacher als je zuvor

TorchNeuron macht die Ausführung von Hugging Face-Modellen auf AWS Trainium einfacher als je zuvor
2 Minuten Lesezeit

Hugging Face kündigt native Unterstützung für Trainium, den AWS-Chip, durch TorchNeuron an, ein neues Backend für PyTorch. Transformers, Diffusers und Kernels können nun direkt auf diesem Chip ausgeführt werden, was die Ausführung von Modellen auf Trainium und deren Verteilung über den Hugging Face Hub erleichtern soll. Laut Quelle wird ab dem 20. Oktober 2026 eine öffentliche Beta-Version verfügbar sein.

TorchNeuron ist eine Erweiterung für PyTorch, die Referenzbibliothek für die Entwicklung von Machine-Learning-Modellen. Sie fügt PyTorch ein für Trainium dediziertes Backend hinzu, mit dem « Eager »-Ausführungsmodus, torch.compile (ein Tool zur Codeoptimierung) und nativen verteilten APIs wie FSDP und TP. Forscher und Entwickler können es mit dem Befehl pip install torch-neuronx installieren.

Ein Ansatzwechsel für die Hugging Face-Bibliotheken

Bisher erfolgte die Nutzung von Trainium und Inferentia, zwei von AWS entwickelten Chips – einer für kostengünstiges Modelltraining, der andere für skalierbare Inferenz – über die Bibliothek optimum-neuron. Diese basierte auf dem Paket torch-neuronx, das selbst auf PyTorch XLA aufbaute, was ein Denken « wie ein Compiler » erforderte: träge Tensoren, ausschließlich statische Formen und Rekompilierung bei jedem Graphbruch. Hugging Face gibt an, dass optimum-neuron dazu diente, diese Komplexität durch speziell gepflegten Code zu verbergen.

TorchNeuron ändert diesen Ansatz, indem es sich direkt in PyTorch integriert, das Neuron nun als natives Gerät erkennt. Trainium wird damit ein Beschleuniger wie jeder andere für Transformers, Diffusers, Accelerate, PEFT und TRL, mit denselben Klassen und Skripten. Hugging Face empfiehlt Neuron-Nutzern nun, auf diese Bibliotheken statt auf Optimum Neuron zurückzugreifen.

Für die Textgenerierung ist die Unterstützung umfangreich: Laut Hugging Face sind 89 % der 100 meistgenutzten Textgenerierungsmodellarchitekturen sofort kompatibel, gegenüber 46 % mit Optimum Neuron v0.4.5. Es genügt, das Argument device_map=neuron zu übergeben, um das Modell auf Neuron-Geräte zu verschieben.

Training, Diffusion und Rechenkerne

Für das Training funktioniert ein Standard-Fine-Tuning-Skript mit TRL auf einer Trainium-Instanz mit einigen Neuron-spezifischen Anpassungen. Hugging Face präzisiert, dass das veröffentlichte Beispiel Qwen/Qwen3-8B und den Datensatz HuggingFaceH4/ultrachat_200k verwendet und auf 8 Trainium2-Chips, also 32 NeuronCores, ausgeführt wurde – das entspricht der Hälfte einer trn2.48xlarge-Instanz. Der erzeugte Sicherungspunkt ist ein standardmäßiger Transformers-Checkpoint, der auf anderer Hardware verwendbar ist.

Mehrere Pipelines von Diffusers sind mit TorchNeuron kompatibel, sowohl im « Eager »- als auch im kompilierten Modus. FLUX.1-dev, FLUX.2, Qwen-Image, Qwen-Image-2.1 und H3 unterstützen auch Tensor-Parallelismus, der jede Schicht des Modells auf mehrere Kerne verteilt. Hugging Face gibt an, dass weitere Pipelines schrittweise kompatibel gemacht werden und ein Leitfaden das Hinzufügen dieser Unterstützung ermöglicht.

Schließlich wird die Kernels-Bibliothek für Neuron optimierte Rechenkerne aufnehmen, die mit der Neuron Kernel Interface (NKI) geschrieben sind. Ein für Vision-Transformer veröffentlichter Kernel ermöglicht laut Quelle auf Trainium eine Beschleunigung « bis zu x2,4 ». Hugging Face dankt dem AWS-Team für seine Arbeit an TorchNeuron und seine Reaktionsfähigkeit auf das Feedback der Community.

S’abonner
Notification pour
guest

0 Commentaires
Le plus ancien
Le plus récent Le plus populaire
0
Nous aimerions avoir votre avis, veuillez laisser un commentaire.x