Die Versionshinweise zu AWS Neuron 2.33 datieren die öffentliche Beta von TorchNeuron auf den 19. Oktober. Die PyTorch-Ankündigung, ebenfalls vom 19. Oktober, nennt den 20. Oktober als Start der öffentlichen Beta und der Open-Source-Veröffentlichung. Das Paket torch-neuronx registriert für Trainium2 und Trainium3 ein neuron Gerät bei PyTorch. Dadurch lassen sich für unterstützte Modell- und Tensoroperationen vertraute PyTorch-APIs verwenden. Für Ingenieure, die diese Chips erwägen, ist der nützliche Unterschied der kürzere Weg von einem bestehenden PyTorch-Programm zu einem Versuch auf dem Gerät. Ob das Programm korrekt und effizient läuft, hängt weiterhin von seinen Operationen, Formen und seiner Softwareumgebung ab.

Der bisherige PyTorch-Weg für Trainium nutzte eine XLA-basierte torch-neuronx Bibliothek und torch_xla Gerätecode. Laut AWS verwendet das native Backend die PyTorch-Erweiterung PrivateUse1 und ersetzt xm.xla_device() durch die Platzierung auf dem Gerät neuron. Für Projekte, die Neuron bereits einsetzen, ist das eine wesentliche Schnittstellenänderung, auch wenn der Paketname vertraut ist. Die Version gilt für Training und Inferenz auf Trn2- und Trn3-Instanzen; die Unterstützungsmatrix führt Trn1 nicht auf.

Was sich in einem PyTorch-Programm ändert

Die AWS-Versionshinweise zeigen ein Modell, das mit .to("neuron") verschoben wird, und einen Eingabetensor, der mit device="neuron" erstellt wird. Unterstützte Operationen mit torch.nn, Autograd und Optimierern laufen anschließend über das Backend. Code für die geräteunabhängige torch.accelerator API von PyTorch kann Neuron ohne eine gerätespezifische Neufassung erkennen, so die TorchNeuron-Übersicht.

Es gibt zwei Ausführungspfade. Der Eager-Modus führt PyTorch-Operationen aus, wenn Python sie aufruft, und erhält den üblichen Kontrollfluss, Ausgaben und Haltepunkte. Jede Operation wird bei der ersten Verwendung kompiliert; spätere passende Aufrufe können einen Cache verwenden. Ingenieure können auch ein Modell oder einen ausgewählten Bereich mit torch.compile(model, backend="neuron", dynamic=False) kompilieren. Das native Backend kompiliert unterstützte Vorwärts- und Rückwärtsgraphen; Code außerhalb des erfassten Bereichs läuft weiter im Eager-Modus. Die PyTorch-Ankündigung beschreibt außerdem, wie man vor dem Kompilieren einmal torch.compiler.set_default_backend("neuron") setzt, damit spätere Kompilierungsaufrufe das Backend-Argument weglassen können.

Die Version enthält ein neuron Backend für torch.distributed mit dokumentierten kollektiven und Punkt-zu-Punkt-Operationen. AWS nennt DistributedDataParallel, FSDP2, DTensor und DeviceMesh sowie Pipeline-Parallelität unter den unterstützten APIs. Das sind nützliche Ansätze, die man für ein Projekt zum verteilten Training prüfen kann. Die Versionshinweise belegen jedoch nicht, dass ein bestehender Job mit mehreren Knoten unverändert läuft oder ein bestimmtes Durchsatzziel erreicht.

Die Evaluation beginnt mit einer passenden Umgebung

Die Installations-Kurzanleitung von AWS erfordert eine Trn2- oder Trn3-Instanz sowie entweder eine Neuron Base DLAMI oder eine unterstützte Ubuntu AMI. Die Versionsmatrix nennt Linux x86_64 unter Ubuntu 24.04 oder 26.04. Für Ubuntu 22.04 endet die Unterstützung mit dieser Version. Aufgeführt sind Python 3.12, 3.13 und 3.14 sowie torch-neuronx 2.13.3.0 für PyTorch 2.13.0 oder torch-neuronx 2.14.3.0 für PyTorch 2.14.0. Die Paketfamilie muss zur PyTorch-Version passen.

Die Kurzanleitung installiert torch-neuronx==2.14.* zusammen mit neuronx-cc==2.* und nki aus den Paketindizes von Neuron und PyTorch. Anschließend prüft sie die Verfügbarkeit des Geräts und vergleicht eine kleine Matrixoperation auf Neuron mit einem CPU-Ergebnis. Damit wird der installierte Gerätepfad im AWS-Beispiel überprüft, nicht das Modell des Lesers. Eine sinnvolle nächste Evaluation ist, die Vorwärts- und Rückwärtspfade des tatsächlichen Modells auszuführen, aufgezeichnete CPU-Fallbacks mit torch_neuronx.get_fallback_ops() zu untersuchen, Ausgaben mit der bisherigen Umgebung zu vergleichen und Kompilierung sowie Ausführungen im eingeschwungenen Zustand getrennt zu messen. Diese letzten Schritte sind von BIG CHANGE vorgeschlagene Evaluationsfragen, keine von uns durchgeführten Tests.

Wo die Beta die Arbeit noch verändern kann

Im Eager-Modus kann eine nicht unterstützte Operation auf die CPU zurückfallen und eine Warnung ausgeben. Unter torch.compile scheitert die Kompilierung standardmäßig, wenn eine im Graph erfasste Operation vom Neuron-Backend nicht in eine niedrigere Ausführungsdarstellung überführt werden kann. Für die Abdeckung einzelner Operationen verweist AWS Entwickler auf seine Kompatibilitätstabelle für Operatoren. Ein abgeschlossener Eager-Lauf belegt daher für sich genommen weder die vollständige Ausführung auf dem Gerät noch die Kompatibilität mit kompilierter Ausführung.

Die Versionshinweise sagen außerdem, dass torch.compile die Einstellung dynamic=False erfordert; eine neue Eingabeform kann eine weitere Kompilierung nötig machen. AWS empfiehlt Padding oder Bucketing, um die Formenvielfalt zu verringern. Der erste Lauf umfasst die Kompilierung, wenn kein passendes Artefakt im Cache liegt. Bei Zeitmessungen eines Workloads sind diese Details wichtig: Die Kompilierungslatenz und das Muster der Eingabeformen können das Evaluationsergebnis verändern.

Weitere ausdrückliche Grenzen schränken die Kompatibilitätsaussage ein. Die Beta unterstützt weder torch.sparse noch torch.func noch torch.nested noch torch.distributions. Nur einige torch.fft Transformationen laufen auf Neuron, und float64 wird als float32 ausgeführt, wobei eine Warnung erscheint. In dieser Version fehlen die TorchAO-Integration und Unterstützung für den zusammengesetzten Datentyp MXFP8. Ein Team, das diese Schnittstellen nutzt, muss den eigenen Codepfad prüfen, bevor es die Geräteplatzierung als Migrationsplan versteht.

Die große Veränderung

TorchNeuron stellt Trainium in einer öffentlichen Beta hinter die nativen Geräte- und Kompilierungsschnittstellen von PyTorch. Das macht einen ersten Versuch mit unterstützten PyTorch-Programmen direkter und bietet bisherigen Neuron-Nutzern einen anderen Weg als ihren XLA-basierten Code. Die entscheidenden Belege für eine Einführung muss jeder Workload liefern: Bleiben seine Operatoren auf dem vorgesehenen Gerät, sind seine numerischen Ergebnisse weiterhin akzeptabel und rechtfertigen die gemessene Leistung und die Kosten die Instanz sowie den Migrationsaufwand? BIG CHANGE hat für diese Analyse Dokumentation und Ankündigungen geprüft; wir haben TorchNeuron nicht ausgeführt oder die Leistungsangaben von AWS unabhängig bestätigt.

Quellen und weiterführende Lektüre