Läuft Ihr Modell auf AWS Trainium? Schlagen Sie es nach

In einem früheren Beitraghaben wir TorchNeuron vorgestellt, das native PyTorch-Backend für AWS Trainium (das torch-neuronx -Paket). Damit führt Transformers seinen eigenen Modellierungscode auf Trainium aus, wie auf jedem anderen Gerät auch, sodass jedes Transformers-Modell auf dem Hub laufen kann. Welche tun das tatsächlich?

Der torch-neuronx model support -Space beantwortet diese Frage für die beliebtesten Modelle auf dem Hub. Er sagt Ihnen, ob ein Modell unverändert auf Trainium läuft, und liefert Ihnen das Skript, um es selbst auszuführen.

Ein Modell nachschlagen

Der Space öffnet mit der Selected -Ansicht: unsere eigene Auswahl von Modellen, getroffen aus den Top-Modellen auf dem Hub nach Downloads und nach Trending-Score. Jede Zeile gibt das Ergebnis an, die Anzahl der Trainium-Cores, über die das Modell geshardet ist, die Wall Time des Laufs (einschließlich Kompilierung) und den Spitzenwert des Gerätespeichers.

Um ein bestimmtes Modell nachzuschlagen, wechseln Sie zur Trending - oder Downloads -Ansicht und tippen Sie einen Teil seines Namens ein. Diese Ansichten folgen den eigenen Ranglisten des Hubs, jeweils eine Aufgabe gleichzeitig, unabhängig davon, ob ein Modell getestet wurde oder nicht. Wenn ein Modell absichtlich ausgelassen wurde, sagt Ihnen der Space, warum (meistens wurde es nicht priorisiert, manchmal ist es wirklich inkompatibel). Sie können auch die Versionen des Neuron-Compilers (neuronx-cc) und von Transformers auswählen, die Sie betrachten möchten.

Zum Zeitpunkt der Erstellung dieses Texts werden nur Transformers-Modelle getestet, und unsere Auswahl ist auf text-generation - und image-text-to-text -Modelle beschränkt. Wir planen, auf diffusers, sentence transformers und andere Transformers-Aufgaben zu erweitern.

Selbst ausführen

Wählen Sie ein Modell aus, um zu sehen, wie man es ausführt. Der Space listet den Neuron-Stack auf, auf dem das Modell getestet wurde, die zu installierenden Pakete und das run.py -Skript, das ausgeführt wurde. Das Skript folgt dem Rezept der TorchNeuron announcement: es lädt das Modell mit from_pretrained, geshardet, wenn es nötig ist, und generiert dann mit einem statischen Cache und torch.compile(backend="neuron").

Auf einer Trainium2 (trn2-)Instanz mit diesem Stack, zum Beispiel im AWS Neuron Deep Learning Container, installieren Sie die Pakete und führen Sie dann aus:

python3 run.py

Kein Bedarf an torchrun: das Skript startet sein eigenes torch.distributed Ranks, wenn das Modell geshardet ist.

Ein erster Lauf beginnt damit, die Graphen des Modells zu kompilieren. Um diesen Schritt zu überspringen, kopieren Sie die Graphen, die derselbe Compiler während des Tests erstellt hat, aus dem öffentlichen NEFF-Cache -Bucket. Der Space liefert Ihnen den Befehl mit der richtigen Compiler-Version:

hf buckets sync hf://buckets/aws-neuron/torch-neuronx-neff-cache/neuronxcc-<version> /tmp/neff_cache

Was „läuft“ bedeutet

Jeder Test lädt das Modell auf einer trn2 -Instanz, geshardet über die kleinste Anzahl von Trainium-Cores, die seine Gewichte fassen kann. Die meisten Modelle werden mit Tensor-Parallelismus geshardet. Mixture-of-experts-Modelle ohne Tensor-Parallelitätsplan verwenden ausschließlich Expert-Parallelismus. Der Test erzeugt dann 32 Tokens und prüft sie. Auf einem einzelnen Core werden die Tokens mit denen eines CPU-Laufs desselben Checkpoints verglichen. Wenn das Modell geshardet ist, prüft ein kleines Judge-Modell, ob der Text flüssiges Englisch ist.

Zum Zeitpunkt der Erstellung dieses Texts werden nur Modelle getestet, die auf 4 Trainium2-Geräten passen, was höchstens 16 Cores bedeutet. Ein PASS sagt nichts über den Durchsatz oder über längere Generierungen aus.

Die Ergebnisse sind ein Datensatz

Alles, was der Space zeigt, stammt aus dem öffentlichen aws-neuron/torch-neuronx-model-support -Datensatz, mit einer Parquet-Datei pro Paar aus neuronx-cc - und Bibliotheksversionen.

Jede Zeile enthält das Ergebnis eines Modells, den Fehler bei einem Fehlschlag, die Versionen des gesamten Stacks und den run.py , der ausgeführt wurde. Sie können ihn direkt abfragen:

import pandas as pd

df = pd.read_parquet("hf://datasets/aws-neuron/torch-neuronx-model-support/data")
df[df.model_id == "Qwen/Qwen3-0.6B"][["neuronx_cc", "library", "library_version", "neuron_cores", "result"]]

Die Datensatzkarte erklärt die genaue Bedeutung jedes Ergebnisses.

Wenn Sie sich das nächste Mal fragen, ob ein Modell auf Trainium läuft, schlagen Sie es nach.