AWS öffnet TorchNeuron, ein natives PyTorch-Backend für Trainium, in der öffentlichen Beta

AWS öffnet TorchNeuron, ein natives PyTorch-Backend für Trainium, in der öffentlichen Beta

AWS gab am 20. Oktober bekannt, dass TorchNeuron, ein quelloffenes PyTorch-Backend für die KI-Chips Trainium, in die öffentliche Beta gegangen ist. Ein längerer Beitrag von AWS-Ingenieuren im PyTorch-Blog trägt das Datum 19. Oktober und wurde am 20. Oktober aktualisiert, dem Tag der PyTorch Conference North America. Die Ankündigung von AWS sagt, der Code sei für Amazon-EC2-Instanzen mit Trainium2 und Trainium3 quelloffen. Der Beitrag nennt Apache 2.0 als Lizenz und ein GitHub-Repository unter aws-neuron. Diese Ausgabe vom 27. Oktober berichtet eine Woche nach dem ursprünglichen Datum. Beide Seiten geben die Sicht von AWS wieder; der PyTorch-Beitrag zitiert jedoch den Technikchef von Deepgram: In den vom Unternehmen ausgewerteten Benchmarks verarbeitete Trainium pro Dollar etwa 60 % mehr Sekunden eingehender Audiodaten als dessen bisherige GPU-Basis. Das ist eine Kundenaussage, kein unabhängiger Test und keine allgemeine Leistungszusage für TorchNeuron.

Was sich für Entwickler ändert

Bisher bedeutete PyTorch-Code auf Trainium den Einsatz getrennter AWS-Werkzeuge. TorchNeuron meldet Trainium über PrivateUse1, PyTorchs offiziellen Erweiterungspunkt für eigene Hardware, und über die Schnittstelle `torch.accelerator` als gewöhnliches PyTorch-Gerät an. Im einfachsten Fall verschiebt ein Entwickler ein Modell mit `.to('neuron')` und behält den Rest der Trainingsschleife bei. Zwei Ausführungsmodi werden unterstützt. Der Eager-Modus führt Operationen einzeln aus, was sich zur Fehlersuche eignet, und der Beitrag fügt einen experimentellen „Adaptive Eager“-Modus hinzu, der die Startkosten senken und Operationen ohne Codeänderung zusammenfassen soll. Der Modus `torch.compile` erfasst das Modell als Graph und übergibt es dem Neuron-Compiler; AWS sieht ihn für die Produktion vor. Der Beitrag nennt den Kompilierungspfad: TorchDynamo, AOTAutograd, Torch-MLIR, StableHLO und dann der Neuron-Compiler, der eine auf dem Chip ausgeführte Datei erzeugt.

Umfang und Grenzen

AWS nennt eine breite Operatorabdeckung für Text-, Bild-, Video- und andere Modelle, Unterstützung der Typen BF16, FP32 und FP16 sowie die üblichen verteilten APIs: DDP, FSDP2, Tensorparallelität, Pipeline-Parallelität, DTensor und DeviceMesh. Die Kommunikation läuft innerhalb eines Knotens über Trainiums Kollektiv-Kerne und zwischen Knoten über den Elastic Fabric Adapter. Der Beitrag spricht Lücken offen an. TorchAO und FP8 befinden sich in einer privaten Beta. Nicht-Standard-Streams teilen sich eine einzige Hardware-Warteschlange, sie behalten also die Reihenfolge, erhalten aber keine eigene Engine. Fehler zeigen sich am nächsten Synchronisationspunkt statt an der fehlerhaften Operation, und TorchNeuron wirft den frühesten mit Kontext erneut. Eigene Kernel werden mit dem Neuron Kernel Interface geschrieben, und AWS liefert eine quelloffene Bibliothek solcher Kernel, darunter Routinen für Attention und Mixture-of-Experts. Eine Profiler-Anbindung verknüpft Framework-Ereignisse mit der Hardwareaktivität.

Wie man die Ankündigung liest

Im Mittelpunkt der Produktbeschreibung steht die PyTorch-Kompatibilität. AWS sagt, Entwickler behielten die PyTorch-APIs und bekämen zugleich Zugang zu Trainium, und ein Engineering-Direktor von Meta wird zitiert, der die Unterstützung begrüßt. Das Ziel des Adaptive-Eager-Modus, die Leistung kompilierter Graphen anzunähern, ist ein im Beitrag genanntes Ziel. Der PyTorch-Beitrag enthält die genannte Zahl von Deepgram zur Kosteneffizienz bei eingehenden Audiodaten, aber weder eine unabhängige Prüfung noch eine allgemeine Leistungszusage. Die dort beschriebenen 100.000 Modellläufe prüfen Hugging-Face-Modelle in reduzierten, repräsentativen Konfigurationen auf Laden, Ausführung, Kompilierung und Operatorplatzierung; sie sind keine 100.000 Leistungsvergleiche mit GPUs. Ein Team, das Trainium erwägt, sollte sein eigenes Modell im Eager- und im kompilierten Modus laufen lassen, prüfen, ob die benötigten Operatoren abgedeckt sind, und die Hinweise zur Beta lesen, bevor es Produktionsarbeit verlagert.