Os chips Trainium da Amazon há tempos são, no papel, um lugar mais barato para treinar modelos e, na prática, um lugar incômodo, porque o código PyTorch precisava ser adaptado para rodar neles. O TorchNeuron, um backend nativo do PyTorch para o Trainium, quer acabar com esse atrito. Entrou em beta pública em 20 de outubro de 2026, instala com pip install torch-neuronx e, segundo o Hugging Face, Transformers, Diffusers, Accelerate, PEFT e TRL já rodam diretamente nele.
Até agora, o caminho do Hugging Face para o Trainium e o Inferentia, os chips de treino e de inferência da AWS, passava pelo Optimum Neuron, construído sobre um pacote do PyTorch/XLA. Isso exigia pensar como um compilador: tensores preguiçosos em vez de execução imediata, só formatos estáticos e uma recompilação sempre que o grafo mudava. O TorchNeuron expõe o chip como um dispositivo neuron comum, com execução eager, torch.compile e as ferramentas distribuídas do próprio PyTorch, como FSDP2 e DTensor. O Hugging Face agora recomenda usar suas bibliotecas diretamente, em vez de passar pelo Optimum Neuron.
Os números de cobertura são a parte útil. Com o TorchNeuron, 89% das 100 arquiteturas de geração de texto mais usadas funcionam de cara, contra 46% com o Optimum Neuron v0.4.5. O ajuste fino usa um script padrão do TRL com algumas configurações específicas do Neuron, e o checkpoint resultante é um arquivo comum do Transformers, que pode ir para o Hub e ser carregado em qualquer hardware. No exemplo detalhado, 32 processos numa instância dividem cada modelo entre quatro núcleos com paralelismo de tensores e distribuem o restante com FSDP. No Diffusers, modelos de imagem como FLUX.1-dev, FLUX.2 e Qwen-Image-2.1 também aceitam paralelismo de tensores, e é possível escrever kernels próprios na Neuron Kernel Interface da AWS e compartilhá-los no Hub.
Junto com o TorchTPU para os chips do Google na mesma semana, é uma mudança notável para quem orça treinamento fora da Nvidia. Os dois backends fazem esses aceleradores parecerem só mais um dispositivo PyTorch para as bibliotecas que analisamos, entre elas PyTorch e Transformers, o que barateia o teste.
É uma beta, e as arestas aparecem. Os formatos ainda precisam ser estáticos para evitar recompilações, então os scripts de treino desligam os formatos dinâmicos automáticos e completam os lotes; o carregamento assíncrono de pesos do Transformers ainda não funciona; e cobertura de 89% significa que algumas arquiteturas populares ainda não rodam. O texto não traz comparações de velocidade ou custo com GPUs, que é o número de que a maioria das equipes vai precisar. Teste uma carga representativa antes de planejar a mudança.