Beta do TorchNeuron leva PyTorch nativo ao AWS Trainium

O TorchNeuron, uma nova forma de executar código PyTorch nos chips Trainium da Amazon, entrou em beta pública em 20 de outubro de 2026. Isso importa porque a maioria dos modelos abertos populares agora roda no Trainium sem alterações. Segundo o Hugging Face, 89% das 100 arquiteturas de modelos de geração de texto mais usadas funcionam de imediato, ante 46% com as ferramentas antigas.

O que é o TorchNeuron

O Trainium é o chip de IA que a Amazon projeta e aluga por meio da AWS, seu negócio de nuvem. O PyTorch é a biblioteca mais usada para criar e treinar modelos de IA. Até agora, o código PyTorch chegava ao Trainium por meio de uma camada de tradução, e muitos modelos precisavam de trabalho extra para rodar.

O TorchNeuron é um plugin do PyTorch que adiciona o Trainium como dispositivo nativo. A documentação da AWS diz que ele “registra o Trainium como um dispositivo padrão do PyTorch”. Na prática, um desenvolvedor move um modelo para o dispositivo chamado neuron, da mesma forma que o moveria para uma GPU com cuda.

O projeto é de código aberto sob a licença Apache-2.0, segundo sua página no GitHub. A AWS o lista como beta pública na versão Neuron 2.33, seu kit de software para os chips.

O que mudou em relação ao caminho antigo

A rota antiga passava pelo Optimum Neuron, uma biblioteca do Hugging Face, que dependia do PyTorch XLA. O XLA é uma camada de compilação que transforma um modelo em um grafo fixo antes da execução. Essa abordagem dificultava o suporte a muitos modelos, porque cada um precisava de sua própria adaptação.

O TorchNeuron elimina essa camada. As próprias bibliotecas do Hugging Face agora rodam diretamente no Trainium: Transformers para modelos de linguagem, Diffusers para modelos de imagem, além de Accelerate, PEFT e TRL para treinamento e ajuste fino. A tabela abaixo resume os recursos, com base na documentação da AWS e no post do Hugging Face.

Recurso	O que significa
Modo eager	O código roda linha por linha, como no PyTorch comum
torch.compile	Compilação opcional para ganhar velocidade, com o backend neuron
Treinamento distribuído	FSDP, DDP, DTensor, paralelismo de tensores e de pipeline
Kernels personalizados	Código escrito à mão para o chip por meio da linguagem NKI da AWS
Chips	Trainium 2 e Trainium 3; FP8 e MXFP8 exigem o Trainium 3

FSDP e DDP são formas padrão do PyTorch de dividir o treinamento entre muitos chips. FP8 é um formato numérico menor que economiza memória e tempo.

Modelos e velocidade

O post do Hugging Face mostra um ajuste fino do Qwen3-8B com TRL, dividido entre vários chips com paralelismo de tensores e FSDP. Ele também executa pipelines de imagem como FLUX.1-dev, FLUX.2 e Qwen-Image com paralelismo de tensores.

A documentação da AWS lista modelos de mistura de especialistas (mixture-of-experts), que encaminham cada entrada para algumas sub-redes especializadas. DeepSeek, GPT-OSS e Qwen3 são suportados. O treinamento desses modelos passa por uma ferramenta separada, o TorchTitan Neuron.

O único número de velocidade tem alcance restrito. Um kernel de vision transformer da equipe Neuron da AWS oferece “uma aceleração de até 2,4 vezes no Trainium”, diz o post do Hugging Face. O post não diz em relação a quê isso é medido.

Como instalar

A página no GitHub traz uma instalação com versões fixadas, executada depois que o driver e os pacotes de runtime do Neuron estiverem instalados:

pip install torch==2.14.0 --index-url https://download.pytorch.org/whl/cpu
pip install "torch-neuronx==2.14.*" "neuronx-cc==2.*" "nki>=0.6" --extra-index-url https://pip.repos.neuron.amazonaws.com

As versões fixadas importam. O TorchNeuron é distribuído com o mesmo nome de pacote, torch-neuronx, que o produto antigo baseado em XLA. Versões anteriores instalam o antigo, alerta o repositório. Executar python -m torch_neuronx.collect_env confirma que o dispositivo pode ser usado a partir do PyTorch. As imagens de máquina e os contêineres prontos da AWS já vêm com ele instalado.

O que isso significa para desenvolvedores

O TorchNeuron transforma o Trainium de um alvo especial em mais um dispositivo do PyTorch. Esse é o passo que torna práticas as comparações de preço com GPUs para equipes comuns.

• Teste primeiro o seu próprio modelo. Com 89% de cobertura, ainda resta cerca de uma em cada dez arquiteturas populares que precisa de trabalho. Rode um job curto de treinamento ou inferência antes de planejar uma migração.
• Fixe as versões em todos os ambientes. Um torch-neuronx sem versão fixada em um arquivo de requirements pode instalar o produto XLA antigo. Verifique a saída do collect_env na CI.
• Confira a geração do chip. Só instâncias com Trainium 2 e 3 são suportadas. Se o seu plano depende de treinamento em FP8, você precisa do Trainium 3.
• Trate-o como beta. Nenhuma das três fontes diz que ele está pronto para produção. Mantenha uma alternativa em GPU para tudo o que tiver prazo.
• Compare o custo por job, não por hora. O número de 2,4 vezes cobre um único kernel, sem base de comparação informada. Meça seus próprios tokens por dólar antes de comprometer orçamento.