Os modelos da biblioteca Transformers podem funcionar nos chips Trainium da AWS graças ao TorchNeuron, um backend nativo para PyTorch, a biblioteca open source de aprendizagem automática. Resta saber quais funcionam realmente. O Space torch-neuronx model support, publicado pela Hugging Face, responde a esta questão para os modelos mais populares do Hub.
Procurar um modelo
O Space abre com a vista Selected, uma seleção de modelos escolhidos entre os mais descarregados e mais populares do Hub. Cada linha indica o resultado do teste, o número de núcleos Trainium nos quais o modelo está distribuído, a duração total da execução, compilação incluída, e a memória máxima utilizada no dispositivo.
Para procurar um modelo específico, basta passar para as vistas Trending ou Downloads e introduzir uma parte do seu nome. Estas vistas seguem a classificação do Hub, tarefa por tarefa, quer o modelo tenha sido testado ou não. Quando um modelo foi excluído, o Space especifica a razão: na maioria das vezes, não era prioritário, por vezes é realmente incompatível. O utilizador também pode escolher as versões do compilador Neuron (neuronx-cc) e do Transformers a examinar.
Por enquanto, apenas os modelos Transformers são testados, e a seleção limita-se às tarefas de geração de texto e de texto-imagem para texto. A equipa planeia estender o dispositivo a diffusers, aos sentence transformers e às outras tarefas do Transformers.
Executar um modelo por si mesmo
Ao selecionar um modelo, vê-se a pilha de software Neuron na qual foi testado, os pacotes a instalar e o script run.py utilizado durante o teste. Este script carrega o modelo com from_pretrained, distribuído por vários núcleos se necessário, e depois gera texto com uma cache estática e torch.compile(backend=«neuron»).
Numa instância trn2 dotada da mesma pilha, por exemplo no AWS Neuron Deep Learning Container, basta instalar os pacotes e depois executar python3 run.py. Não é necessário nenhum torchrun, o script lança ele próprio os processos de distribuição quando o modelo está distribuído. A primeira execução começa pela compilação dos grafos do modelo. Para a dispensar, pode-se copiar os grafos produzidos durante o teste a partir do bucket público NEFF cache, com o comando fornecido pelo Space, que tem em conta a versão do compilador.
O que significa «funciona»
Cada teste carrega o modelo numa instância trn2, distribuído pelo menor número de núcleos Trainium capaz de conter os seus pesos. A maioria dos modelos usa o paralelismo de tensores, que divide os cálculos entre vários núcleos. Os modelos de especialistas mistos, sem plano de paralelismo de tensores, usam apenas o paralelismo de especialistas. O teste gera depois 32 tokens, ou seja, unidades de texto, e verifica-os. Num único núcleo, estes tokens são comparados com os de uma execução em CPU do mesmo checkpoint. Quando o modelo está distribuído, um pequeno modelo de juízo verifica que o texto é um inglês fluente.
Apenas os modelos que cabem em quatro dispositivos Trainium2 são testados, ou seja, no máximo dezasseis núcleos. Um resultado PASS não diz nada sobre o débito nem sobre gerações mais longas.
Todos os resultados provêm do conjunto de dados público aws-neuron/torch-neuronx-model-support, que contém um ficheiro Parquet por par de versões neuronx-cc e de biblioteca. Cada linha dá o resultado de um modelo, o eventual erro, as versões da pilha e o run.py executado. É possível consultá-lo diretamente com pandas, por exemplo para o modelo Qwen/Qwen3-0.6B. A ficha do conjunto de dados especifica o significado exato de cada resultado.
