A Hugging Face anuncia o suporte nativo para o Trainium, o chip da AWS, através do TorchNeuron, um novo backend para o PyTorch. Transformers, Diffusers e Kernels podem agora ser executados diretamente neste chip, o que deve facilitar a execução de modelos no Trainium e a sua partilha no Hub da Hugging Face. Segundo a fonte, uma versão beta pública estará disponível a partir de 20 de outubro de 2026.
TorchNeuron é uma extensão do PyTorch, a biblioteca de referência para o desenvolvimento de modelos de aprendizagem automática. Adiciona ao PyTorch um backend dedicado ao Trainium, com o modo de execução « eager », torch.compile (uma ferramenta de otimização de código) e as APIs distribuídas nativas como FSDP e TP. Investigadores e programadores poderão instalá-lo com o comando pip install torch-neuronx.
Uma mudança de abordagem para as bibliotecas Hugging Face
Até agora, a utilização do Trainium e do Inferentia, dois chips concebidos pela AWS, um para o treino económico de modelos e outro para a inferência em larga escala, passava pela biblioteca optimum-neuron. Esta baseava-se no pacote torch-neuronx, que por sua vez se baseava no PyTorch XLA, o que obrigava a pensar « como um compilador »: tensores preguiçosos, formas estáticas apenas e recompilação em cada rutura de grafo. A Hugging Face indica que o optimum-neuron servia para ocultar esta complexidade através de código mantido especificamente.
O TorchNeuron modifica esta abordagem ao integrar-se diretamente no PyTorch, que vê o Neuron como um periférico nativo. O Trainium torna-se então um acelerador como outro qualquer para Transformers, Diffusers, Accelerate, PEFT e TRL, com as mesmas classes e scripts. A Hugging Face recomenda agora aos utilizadores do Neuron que passem por estas bibliotecas em vez do Optimum Neuron.
Para a geração de texto, o suporte é amplo: segundo a Hugging Face, 89 % das 100 arquiteturas de modelos de geração de texto mais utilizadas são compatíveis de imediato, contra 46 % com o Optimum Neuron v0.4.5. Basta passar o argumento device_map=neuron para mover o modelo para os dispositivos Neuron.
Treino, difusão e núcleos de cálculo
Para o treino, um script padrão de fine-tuning com TRL funciona numa instância Trainium mediante alguns ajustes próprios do Neuron. A Hugging Face precisa que o exemplo publicado utiliza o Qwen/Qwen3-8B e o conjunto de dados HuggingFaceH4/ultrachat_200k, e que foi executado em 8 chips Trainium2, ou seja, 32 NeuronCores, isto é, metade de uma instância trn2.48xlarge. O ponto de salvaguarda produzido é um checkpoint Transformers padrão, utilizável noutros hardwares.
Vários pipelines de Diffusers são compatíveis com o TorchNeuron, tanto em modo « eager » como compilado. FLUX.1-dev, FLUX.2, Qwen-Image, Qwen-Image-2.1 e H3 suportam também o paralelismo de tensores, que distribui cada camada do modelo por vários núcleos. A Hugging Face indica que outros pipelines serão progressivamente compatíveis e que um guia permite adicionar este suporte por si mesmo.
Por fim, a biblioteca Kernels acolherá núcleos de cálculo otimizados para o Neuron, escritos com a Neuron Kernel Interface (NKI). Um núcleo destinado aos transformers de visão, publicado no Hub, permite segundo a fonte uma aceleração « até x2.4 » no Trainium. A Hugging Face agradece à equipa da AWS pelo seu trabalho no TorchNeuron e pela sua reatividade face aos retornos da comunidade.
