TorchNeuron bringt native PyTorch-Unterstützung für AWS-Trainium-Chips

Trainium als Standardgerät in PyTorch

Laut der Ankündigung von PyTorch soll TorchNeuron sicherstellen, dass der Modellcode, die Bibliotheken und die Werkzeuge, die Entwickler bereits nutzen, auf AWS Trainium laufen. Trainium wird als Portfolio von KI-Chips beschrieben, die für Training und Inferenz im Cloud-Maßstab ausgelegt sind. Das Backend unterstützt Trainium2 und Trainium3.

Die Integration baut auf beschleunigerunabhängigen Erweiterungspunkten auf, die die PyTorch-Community geschaffen hat, darunter das PrivateUse1-Geräte-Backend und die torch.accelerator-APIs. TorchNeuron registriert Trainium über PrivateUse1, sodass Standard-APIs wie torch.nn, torch.tensor, torch.autocast, torch.distributed und torch.profiler auf dem Gerät neuron funktionieren.

Der Wechsel der Hardware soll nur eine kleine Änderung erfordern. Im genannten Beispiel wird ein Modell mit .to('neuron') auf den Chip verschoben, der Rest einer Trainingsschleife bleibt unverändert. Code, der gegen torch.accelerator geschrieben wurde, die von PyTorch empfohlene geräteunabhängige API, läuft unverändert. Bei kompilierten Modellen benennen Entwickler das Neuron-Backend entweder einmalig mit torch.compiler.set_default_backend("neuron") oder pro Aufruf mit torch.compile(model, backend="neuron").

Jana van Greunen, Director of PyTorch Engineering bei Meta, sagte in der Ankündigung, „Entwickler sollten auf allen Hardwareplattformen eine einheitliche Erfahrung haben“, und Teams könnten die Fähigkeiten von Trainium nutzen, „ohne ihre Workflows neu erfinden oder das PyTorch-Ökosystem verlassen zu müssen“.

Eager-Modus, torch.compile und verteiltes Training

TorchNeuron unterstützt beide Ausführungsmodi von PyTorch. Der Eager-Modus führt Operationen sofort aus, sodass Python-Kontrollfluss, Ausgaben mit print und Breakpoints ohne Kompilierschritt funktionieren. Der Beitrag empfiehlt ihn für Forschung und Debugging. Weil der Eager-Modus Operationen einzeln abarbeitet, kompiliert jede ihre eigene Datei im Neuron Executable File Format (NEFF), wird auf dem Gerät gestartet und leitet Zwischenergebnisse über den HBM weiter.

Um diesen Aufwand zu senken, führt das Team einen experimentellen adaptiven Eager-Modus ein. Laut dem Beitrag verringert er den Start-Overhead und bietet mehr Möglichkeiten, Operationen zusammenzuführen. So nähert er sich der Leistung eines kompilierten Graphen ohne Codeänderungen an und behält zugleich die interaktive Arbeitsweise bei.

Für den Produktivbetrieb erfasst torch.compile das Modell in einem Graphen, den der Neuron-Compiler optimiert. Die Pipeline nutzt TorchDynamo, um einen FX-Graphen zu erfassen, AOTAutograd, um den Rückwärtsgraphen nachzuverfolgen, und Torch-MLIR, um das Ergebnis in StableHLO zu überführen, bevor der Neuron-Compiler eine NEFF-Datei erzeugt. FlexAttention wird nativ unterstützt. Die Neuron-Inferenzstacks nutzen dieselbe Unterstützung: vLLM Neuron bedient LLMs, vLLM Omni Neuron Diffusionsmodelle für die Bild- und Videoerzeugung.

Bei den Datentypen können Modelle BF16, FP32 und FP16 sowie Ganzzahl- und boolesche Typen verwenden. TorchAO und FP8 befinden sich in einer privaten Beta und sollen mit der nächsten Version allgemein verfügbar werden.

Verteilte Workloads können Standard-PyTorch-APIs nutzen, darunter DDP, FSDP2, SimpleFSDP, Tensor-Parallelität und Pipeline-Parallelität, dazu Kombinationen wie FSDP2 mit Tensor-Parallelität und ZeRO. Unterstützt werden DTensor, DeviceMesh, kollektive APIs wie all_gather, reduce_scatter und all_reduce sowie die P2P-APIs send und recv. Kollektive Operationen laufen innerhalb eines Knotens über die dedizierten Collective-Kerne von Trainium und zwischen Knoten über den Elastic Fabric Adapter (EFA).

Streams, eigene Kernel und Profiling

Torch-Streams und -Events werden über torch_neuronx.Stream und torch_neuronx.Event unterstützt, die die Stream-Semantik von PyTorch auf die Hardware-Queues von Trainium abbilden. Der Beitrag nennt einige Einschränkungen: Nicht-Standard-Streams werden auf eine einzelne Hardware-Queue gebündelt, behalten also Reihenfolge, Priorität und Event-Semantik, erhalten aber keine eigene dedizierte Engine. Fehler treten am nächsten Synchronisationspunkt auf, an dem TorchNeuron die früheste Ausnahme mit ihrem ursprünglichen Kontext erneut auslöst.

Mit dem Neuron Kernel Interface (NKI) können Entwickler eigene Operatoren schreiben, optimieren und sowohl im Eager-Modus als auch in torch.compile ausführen. Forscher können es nutzen, um Ideen zu prototypisieren, die eine Hardwarekontrolle auf niedriger Ebene erfordern, etwa neue Modellarchitekturen. Entwickler können leistungskritische NKI-Kernel direkt aufrufen, um Folgen von Standard-ATen-Operationen zu ersetzen.

TorchNeuron liefert außerdem NKI-lib mit, eine quelloffene Sammlung von auf Trainium optimierten Operatoren für dichte und MoE-Modelle, darunter FlashAttention und blockweise MoE-Kernel. Beim Profiling arbeitet Neuron Explorer mit dem Standard-torch.profiler zusammen, um die Ausführung vom Framework bis zur Hardware nachzuverfolgen und Ereignisse mit Quellcodezeilen zu verknüpfen.

Was das für Entwickler bedeutet

Für Teams, die bereits auf PyTorch setzen, senkt TorchNeuron den Aufwand, Trainium auszuprobieren: Bestehender Code, verteilte Strategien und Profiling-Werkzeuge lassen sich übernehmen, und wo mehr Kontrolle nötig ist, stehen eigene Kernel zur Verfügung. Der Code soll als öffentliche Beta unter github.com/aws-neuron/torch-neuronx veröffentlicht werden, die Ankündigung erfolgte vor dem Open-Source-Start auf der PyTorch Conference North America am 20. Oktober. Einige Funktionen, darunter der adaptive Eager-Modus und TorchAO mit FP8, sind weiterhin experimentell oder befinden sich in einer privaten Beta.