TorchNeuron: Amazon öffnet Trainium für PyTorch mit einer Konfigurationsänderung

Das Ende der erzwungenen Umschreibung

Amazon Web Services hat TorchNeuron vorgestellt, eine Open-Source-Bibliothek unter der Apache-2.0-Lizenz, die es ermöglicht, PyTorch-Code direkt auf den Beschleunigern Trainium2 und Trainium3 auszuführen. Das Hauptziel besteht darin, die technische Hürde zu beseitigen, die Entwickler historisch davon abgehalten hat, ihre Modelle auf alternative Hardware zu migrieren: die Notwendigkeit, komplette Anwendungen umzuschreiben, um sie an eine neue Architektur anzupassen.

Seit über einem Jahrzehnt hat sich das Ökosystem von NVIDIA dank CUDA etabiert, einer Software, die es PyTorch-Anwendungen ermöglicht, nativ auf seinen GPUs zu funktionieren. TorchNeuron greift diesen Wettbewerbsvorteil direkt an, indem es sich in die offiziellen Erweiterungspunkte von PyTorch integriert, insbesondere über das PrivateUse1-Backend. Das bedeutet, dass Entwickler ihren üblichen Arbeitsablauf beibehalten und lediglich das Zielsystem durch eine einzige Konfigurationszeile ändern können.

Technische Architektur und Ausführungsmodi

TorchNeuron ist kein einfacher Wrapper, sondern ein natives Backend, das sich tief in den Compiler von PyTorch integriert. Es unterstützt zwei Hauptausführungsmodi: den Eager-Modus, der ideal für Forschung und Debugging ist, und den torch.compile-Modus, der für die Leistung in der Produktion optimiert ist. Im Eager-Modus wird jede Operation individuell in eine NEFF-Datei (Neuron Executable File Format) kompiliert und auf dem Gerät ausgeführt.

Für leistungskritische Workloads ermöglicht die Verwendung von torch.compile, das Modell als Graph zu erfassen, den der Neuron-Compiler global optimiert. Dieser Prozess nutzt Torch-MLIR, um den Graphen in StableHLO zu konvertieren, bevor der finale Code für die Hardware generiert wird. Darüber hinaus unterstützt die Bibliothek Standard-APIs für verteiltes Training wie FSDP, DDP und DTensor, was das Training großer Modelle ohne Änderung der Parallelisierungslogik erleichtert.

Technische Einschränkungen der ersten Version

Obwohl der Launch bedeutend ist, weist die öffentliche Beta-Version von TorchNeuron konkrete Einschränkungen auf, die Teams vor der Migration kritischer Workloads berücksichtigen sollten. Laut der technischen Dokumentation überspringt die erste Version etwa 22.000 Tests der PyTorch-Suite, hauptsächlich im Zusammenhang mit Operatoren, die noch keine native Implementierung in Neuron haben.

Zu den in diesem ersten Launch nicht unterstützten Funktionen gehören dynamische Eingabeformen (dynamic shapes), verdünnte Tensoren und bestimmte Autograd-Operationen im Forward-Modus. Diese Auslassungen bedeuten, dass nicht sofort jeder PyTorch-Code kompatibel sein wird, und Entwickler müssen ihre spezifischen Modelle validieren, um sicherzustellen, dass sie nicht von diesen ausgeschlossenen Funktionen abhängen.

Strategische Auswirkungen und Wettbewerb mit NVIDIA

Der Schritt von Amazon ist nicht nur technisch, sondern auch strategisch. Indem AWS TorchNeuron öffnet und zum Hauptprojekt von PyTorch beiträgt, um das Framework hardware-agnostischer zu machen, zielt es darauf ab, den Software-Vorteil von NVIDIA zu untergraben. Die Idee ist, dass Entwickler Hardware basierend auf Preis, Leistung und Verfügbarkeit wählen, und nicht aufgrund von erzwungener Code-Kompatibilität.

Colin Brace, VP für Software-Technologie bei AWS Annapurna Labs, betonte, dass das Ziel darin bestehe, den 'Lock-in' zu beseitigen, der Entwicklern eine echte Wahl verwehrte. Mit über 30 Anfangskunden, die die Technologie testen, darunter KI-Stimme-Unternehmen wie Deepgram und Cartesia, zeigt Amazon, dass es möglich ist, Effizienzverbesserungen bei den Kosten zu erzielen, ohne die Modellqualität zu opfern.

So starten Sie mit TorchNeuron

Für Entwickler, die TorchNeuron ausprobieren möchten, ist der Installationsprozess einfach. Die Bibliothek ist auf GitHub im Repository aws-neuron/torch-neuronx verfügbar. Nach der Installation erfolgt der Wechsel zu Trainium durch die Änderung des Zielgeräts des Modells.

Es ist wichtig zu beachten, dass TorchNeuron EC2-Trn2- oder Trn3-Instanzen erfordert. Teams, die bereits PyTorch verwenden, können mit kleinen Modellen beginnen, um die Kompatibilität zu validieren, bevor sie auf große Trainings- oder Inferenz-Workloads skalieren. Die Integration mit Hugging Face Transformers ermöglicht es auch, vortrainierte Modelle direkt auf dem Neuron-Gerät zu laden.

import torch

# Definiere ein einfaches Modell
model = torch.nn.Linear(10, 1)

# Migration zu Trainium (eine Zeile Änderung)
model = model.to('neuron')

# Führe eine Inferenz aus
input_data = torch.randn(1, 10, device='neuron')
output = model(input_data)

print(output)

Weiter entdecken

Korrekturen und Transparenz

Wenn du falsche oder veraltete Angaben findest, sende uns die Seite und die zu prüfende Quelle.