I modelli della libreria Transformers possono funzionare sui chip Trainium di AWS grazie a TorchNeuron, un backend nativo per PyTorch, la libreria open source di machine learning. Resta da sapere quali funzionano realmente. Lo Space torch-neuronx model support, pubblicato da Hugging Face, risponde a questa domanda per i modelli più popolari dell’Hub.
- Hugging Face pubblica uno Space che verifica la compatibilità dei modelli Transformers con AWS Trainium.
- Lo Space mostra i risultati dei test sui modelli più popolari e permette di cercarne uno specifico.
- Il test genera 32 token e li verifica, ma un esito PASS non dice nulla su throughput o generazioni lunghe.
- Per ora sono testati solo i modelli di generazione di testo e testo-immagine verso testo, fino a sedici core.
Cercare un modello
Lo Space si apre sulla vista Selected, una selezione di modelli scelti tra i più scaricati e di tendenza dell’Hub. Ogni riga indica il risultato del test, il numero di core Trainium su cui il modello è distribuito, la durata totale dell’esecuzione, compilazione inclusa, e la memoria massima utilizzata sul dispositivo.
Per cercare un modello preciso, basta passare alle viste Trending o Downloads e inserire una parte del suo nome. Queste viste seguono la classifica dell’Hub, task per task, che il modello sia stato testato o meno. Quando un modello è stato scartato, lo Space precisa il motivo: il più delle volte, non era prioritario, a volte è realmente incompatibile. L’utente può anche scegliere le versioni del compilatore Neuron (neuronx-cc) e di Transformers da esaminare.
Per il momento, solo i modelli Transformers sono testati, e la selezione si limita ai task di generazione di testo e di testo-immagine verso testo. Il team prevede di estendere il dispositivo a diffusers, ai sentence transformers e agli altri task di Transformers.
Lanciare un modello da soli
Selezionando un modello, si vede la pila software Neuron su cui è stato testato, i pacchetti da installare e lo script run.py utilizzato durante il test. Questo script carica il modello con from_pretrained, distribuito su più core se necessario, poi genera testo con una cache statica e torch.compile(backend= »neuron »).
Su un’istanza trn2 dotata della stessa pila, ad esempio nel AWS Neuron Deep Learning Container, basta installare i pacchetti e poi lanciare python3 run.py. Non è necessario alcun torchrun, lo script lancia lui stesso i processi di distribuzione quando il modello è distribuito. La prima esecuzione inizia con la compilazione dei grafi del modello. Per evitarla, si possono copiare i grafi prodotti durante il test dal bucket pubblico NEFF cache, con il comando fornito dallo Space, che tiene conto della versione del compilatore.
Cosa significa « funziona »
Ogni test carica il modello su un’istanza trn2, distribuito sul minor numero di core Trainium in grado di contenere i suoi pesi. La maggior parte dei modelli utilizza il parallelismo dei tensori, che suddivide i calcoli tra più core. I modelli a esperti misti, senza piano di parallelismo dei tensori, utilizzano solo il parallelismo degli esperti. Il test genera poi 32 token, cioè unità di testo, e li verifica. Su un singolo core, questi token sono confrontati con quelli di un’esecuzione su CPU dello stesso checkpoint. Quando il modello è distribuito, un piccolo modello giudice verifica che il testo sia un inglese fluente.
Solo i modelli che stanno su quattro dispositivi Trainium2 sono testati, cioè al massimo sedici core. Un risultato PASS non dice nulla sul throughput né sulle generazioni più lunghe.
Tutti i risultati provengono dal dataset pubblico aws-neuron/torch-neuronx-model-support, che contiene un file Parquet per ogni coppia di versioni neuronx-cc e di libreria. Ogni riga fornisce il risultato di un modello, l’eventuale errore, le versioni della pila e il run.py eseguito. È possibile interrogarlo direttamente con pandas, ad esempio per il modello Qwen/Qwen3-0.6B. La scheda del dataset precisa il significato esatto di ogni risultato.
