Descubra quais Transformers rodam no AWS Trainium em 2026

Inteligência Artificial

Conheça os modelos Transformers compatíveis com AWS Trainium e otimize seus projetos com dados detalhados de desempenho.

Leia também: Programar em linguagem natural aumenta erros: entenda por quê

A equipe do TorchNeuron lançou um Space no Hugging Face que revela, em tempo real, quais modelos Transformers populares rodam hoje no hardware AWS Trainium com o backend torch-neuronx, trazendo também scripts prontos, métricas de desempenho e um dataset público com o resultado de cada teste.

Quais modelos são compatíveis com AWS Trainium na Hugging Face?

O novo Space lista uma seleção de modelos Transformers mais baixados e em alta no Hugging Face Hub e indica, para cada um, se ele roda sem modificações no AWS Trainium, quantos núcleos da placa são usados, o tempo total de execução (incluindo compilação) e o pico de memória na instância.

A interface abre em uma aba com modelos escolhidos pela equipe entre os mais populares do Hub. Há ainda visões específicas por ranking de downloads ou por tendência, em que o usuário digita parte do nome do modelo para descobrir se ele já foi testado no Trainium ou não. Quando um modelo fica de fora de propósito, o Space mostra o motivo — em geral, falta de prioridade de teste ou incompatibilidade real.

Por enquanto, o foco está em modelos de geração de texto e de imagem para texto. A equipe planeja ampliar o suporte para diffusers, sentence transformers e outras tarefas baseadas em Transformers, o que deve cobrir uma parcela ainda maior dos modelos usados hoje em aplicações comerciais e de pesquisa.

Além da indicação de compatibilidade, o Space permite escolher a versão do compilador neuronx-cc e da biblioteca Transformers usada no teste, o que ajuda times que precisam reproduzir resultados em um stack específico de software.

Scripts prontos e execução em instâncias Trainium2

Ao selecionar um modelo, o Space exibe todo o ambiente usado nos testes: versão da pilha Neuron, pacotes adicionais necessários e o script run.py que efetivamente rodou o modelo. Esse script é o coração da proposta.

“O script segue a receita do anúncio do TorchNeuron: carrega o modelo com from_pretrained, particiona quando necessário, e gera com torch.compile(backend=”neuron”)”, explica a equipe TorchNeuron. Na prática, isso significa que o código usa a própria API de Transformers, sem gambiarras específicas de hardware, e delega a otimização ao backend neuron.

Os testes rodam em instâncias Trainium2 (trn2) na nuvem da Amazon, usando o AWS Neuron Deep Learning Container como base. O usuário instala apenas alguns pacotes extras e executa o script. Não há necessidade de chamar ferramentas externas de distribuição: o próprio run.py controla as ranks do torch.distributed quando o modelo é particionado.

Cada modelo é carregado em uma instância trn2 e distribuído pelo menor número de núcleos Trainium que consegue acomodar seus pesos. “Cada teste carrega o modelo em uma instância trn2, particionado no menor número de núcleos Trainium que comporta seus pesos”, diz a equipe. Quando o modelo não cabe em um único núcleo, entra em cena o paralelismo tensorial, espalhando os cálculos entre vários núcleos.

Testes com 32 tokens e até 16 núcleos Trainium

Para validar se um modelo realmente roda no Trainium, o Space segue um protocolo fixo: o sistema gera 32 tokens de saída e verifica se o resultado faz sentido. Em execuções em um único núcleo, os tokens gerados são comparados com a saída do mesmo checkpoint rodando em CPU. Em modelos particionados, entra em campo um modelo “juiz”, que avalia se o texto produzido é inglês fluente.

Os testes utilizam até 4 dispositivos Trainium2 por vez, o que equivale a 16 núcleos ativos no limite. Modelos que excedem essa capacidade ficam, por ora, de fora da lista de compatíveis.

Ao fim, cada caso recebe um resultado como PASS (quando a execução funciona). A equipe deixa claro que esse selo não mede velocidade ou robustez para usos intensivos. “Um PASS não indica throughput nem gerações longas, apenas que o modelo roda”, ressalta a equipe TorchNeuron.

Dataset público com erros, versões e scripts

Além da interface gráfica, o projeto expõe todos os dados dos testes em um dataset público, batizado de aws-neuron/torch-neuronx-model-support. Cada combinação de versão do compilador neuronx-cc e da biblioteca Transformers gera um arquivo Parquet próprio.

Em cada linha, o dataset registra o nome do modelo, o resultado do teste, eventuais erros, as versões completas da pilha de software e o script run.py que foi usado. Para quem integra Trainium a pipelines de MLOps ou precisa auditar ambientes, esses arquivos podem ser consultados diretamente por ferramentas de análise de dados, sem passar pelo Space.

O projeto segue uma filosofia de ciência aberta: as falhas também são registradas, expondo de forma transparente quais combinações de modelo e stack ainda não funcionam bem. Isso ajuda tanto usuários finais quanto a própria equipe TorchNeuron a priorizar correções e otimizações.

Compilação, cache NEFF e ganho de tempo

No primeiro uso em Trainium, o compilador neuronx-cc gera gráficos otimizados do modelo, um processo que consome tempo e recursos. Essa etapa é inevitável em execuções inéditas, mas o Space traz um atalho para quem quer acelerar o caminho.

“Para pular a compilação inicial, copie os gráficos do cache público NEFF com o comando fornecido pelo Space”, orienta a equipe. Os mesmos gráficos produzidos durante os testes são disponibilizados em um bucket público, e o próprio Space entrega o comando pronto com a versão correta do compilador. Ao reutilizar esses artefatos, times de desenvolvimento economizam minutos — ou até horas — em pipelines de teste e produção.

Entenda o contexto

O backend torch-neuronx é a integração nativa entre PyTorch e o AWS Trainium, chip especializado em acelerar tarefas de aprendizado profundo na nuvem da Amazon. Com ele, modelos da biblioteca Transformers podem rodar no Trainium usando o próprio código padrão do framework, o que teoricamente permitiria executar qualquer modelo do Hub nesse hardware.

Na prática, porém, detalhes de implementação, tamanhos de modelos, uso de memória e particularidades da pilha de software fazem com que alguns modelos funcionem de primeira, enquanto outros exigem ajustes. O Space lançado pela equipe TorchNeuron busca reduzir essa incerteza com dados públicos: para cada modelo testado, informa se roda ou não, quanto tempo leva, quanta memória consome e que versões de bibliotecas foram usadas.

A iniciativa reduz a barreira de entrada para quem quer experimentar o Trainium em cargas de trabalho intensivas, como modelos de linguagem grande e sistemas de geração de imagens com texto. Ao mesmo tempo, pressiona fornecedores de soluções que ainda não suportam o chip a atualizar seus produtos para não ficar para trás em custo e desempenho.

Para empresas e pesquisadores, o efeito direto é a possibilidade de planejar migrações com menos tentativa e erro, aproveitando código aberto, scripts reproduzíveis e um repositório transparente de sucessos e falhas. Para a comunidade de IA, o projeto adiciona uma camada de visibilidade rara sobre o que, de fato, funciona hoje em hardware especializado na nuvem.