Modellen uit de Transformers-bibliotheek kunnen draaien op de Trainium-chips van AWS dankzij TorchNeuron, een native backend voor PyTorch, de open-sourcebibliotheek voor machine learning. De vraag is welke modellen daadwerkelijk werken. De Space torch-neuronx model support, gepubliceerd door Hugging Face, geeft hierop een antwoord voor de populairste modellen van de Hub.
- Transformers-modellen draaien op AWS Trainium dankzij TorchNeuron.
- De Space torch-neuronx model support toont welke populaire modellen werken.
- Elk model wordt getest op trn2 en genereert 32 tokens.
- Een PASS zegt niets over snelheid of langere generaties.
Een model zoeken
De Space opent met de weergave Selected, een selectie van modellen gekozen uit de meest gedownloade en trending modellen van de Hub. Elke regel toont het testresultaat, het aantal Trainium-cores waarover het model is verdeeld, de totale uitvoeringstijd inclusief compilatie, en het maximale geheugengebruik op het apparaat.
Om een specifiek model te zoeken, schakelt u eenvoudig over naar de weergaven Trending of Downloads en voert u een deel van de naam in. Deze weergaven volgen de ranglijst van de Hub, taak per taak, ongeacht of het model al is getest of niet. Wanneer een model is uitgesloten, geeft de Space de reden aan: meestal was het geen prioriteit, soms is het echt incompatibel. De gebruiker kan ook kiezen welke versies van de Neuron-compiler (neuronx-cc) en Transformers moeten worden onderzocht.
Momenteel worden alleen Transformers-modellen getest, en de selectie is beperkt tot tekstgeneratie- en tekst-naar-beeld-taken. Het team plant om het systeem uit te breiden naar diffusers, sentence transformers en andere Transformers-taken.
Zelf een model draaien
Door een model te selecteren, ziet u de Neuron-softwarestack waarop het is getest, de te installeren pakketten en het script run.py dat tijdens de test is gebruikt. Dit script laadt het model met from_pretrained, verdeeld over meerdere cores indien nodig, en genereert vervolgens tekst met een statische cache en torch.compile(backend=’neuron’).
Op een trn2-instantie met dezelfde stack, bijvoorbeeld in de AWS Neuron Deep Learning Container, volstaat het om de pakketten te installeren en vervolgens python3 run.py uit te voeren. Er is geen torchrun nodig; het script start zelf de distributieprocessen wanneer het model is verdeeld. De eerste uitvoering begint met het compileren van de modelgrafen. Om dit over te slaan, kunt u de tijdens de test geproduceerde grafen kopiëren vanuit de openbare NEFF-cache-bucket, met het commando dat door de Space wordt geleverd, dat rekening houdt met de compiler-versie.
Wat ‘werkt’ betekent
Elke test laadt het model op een trn2-instantie, verdeeld over het kleinste aantal Trainium-cores dat de gewichten kan bevatten. De meeste modellen gebruiken tensorparallelisme, dat berekeningen verdeelt over meerdere cores. Modellen met mixed experts, zonder tensorparallelisme-plan, gebruiken alleen expertparallelisme. De test genereert vervolgens 32 tokens, dat wil zeggen teksteenheden, en controleert ze. Op een enkele core worden deze tokens vergeleken met die van een CPU-uitvoering van hetzelfde checkpoint. Wanneer het model is verdeeld, controleert een klein beoordelingsmodel of de tekst vloeiend Engels is.
Alleen modellen die passen op vier Trainium2-apparaten worden getest, maximaal zestien cores. Een PASS-resultaat zegt niets over de doorvoersnelheid of langere generaties.
Alle resultaten komen uit de openbare dataset aws-neuron/torch-neuronx-model-support, die een Parquet-bestand bevat per paar neuronx-cc- en bibliotheekversies. Elke regel geeft het resultaat van een model, eventuele fouten, de stackversies en het uitgevoerde run.py. Het is mogelijk om deze dataset rechtstreeks te bevragen met pandas, bijvoorbeeld voor het model Qwen/Qwen3-0.6B. De datasetbeschrijving geeft de exacte betekenis van elk resultaat aan.
