Hugging Face anuncia soporte nativo para Trainium, el chip de AWS, gracias a TorchNeuron, un nuevo backend para PyTorch. Transformers, Diffusers y Kernels ahora pueden ejecutarse directamente en este chip, lo que debería facilitar la ejecución de modelos en Trainium y su compartición en el Hub de Hugging Face. Según la fuente, una versión beta pública estará disponible a partir del 20 de octubre de 2026.
TorchNeuron es una extensión de PyTorch, la biblioteca de referencia para el desarrollo de modelos de aprendizaje automático. Añade a PyTorch un backend dedicado a Trainium, con el modo de ejecución « eager », torch.compile (una herramienta de optimización de código) y las API distribuidas nativas como FSDP y TP. Los investigadores y desarrolladores podrán instalarlo con el comando pip install torch-neuronx.
Un cambio de enfoque para las bibliotecas de Hugging Face
Hasta ahora, el uso de Trainium e Inferentia, dos chips diseñados por AWS, uno para el entrenamiento económico de modelos y otro para la inferencia a gran escala, pasaba por la biblioteca optimum-neuron. Esta se basaba en el paquete torch-neuronx, que a su vez se basaba en PyTorch XLA, lo que obligaba a razonar « como un compilador »: tensores perezosos, formas estáticas únicamente y recompilación en cada ruptura de grafo. Hugging Face indica que optimum-neuron servía para ocultar esta complejidad mediante código mantenido específicamente.
TorchNeuron modifica este enfoque integrándose directamente en PyTorch, que ve a Neuron como un periférico nativo. Trainium se convierte entonces en un acelerador como cualquier otro para Transformers, Diffusers, Accelerate, PEFT y TRL, con las mismas clases y scripts. Hugging Face recomienda ahora a los usuarios de Neuron que utilicen estas bibliotecas en lugar de Optimum Neuron.
Para la generación de texto, el soporte es amplio: según Hugging Face, 89 % de las 100 arquitecturas de modelos de generación de texto más utilizadas son compatibles de inmediato, frente a 46 % con Optimum Neuron v0.4.5. Basta con pasar el argumento device_map=neuron para mover el modelo a los dispositivos Neuron.
Entrenamiento, difusión y núcleos de cálculo
Para el entrenamiento, un script estándar de fine-tuning con TRL funciona en una instancia Trainium con algunos ajustes propios de Neuron. Hugging Face precisa que el ejemplo publicado utiliza Qwen/Qwen3-8B y el conjunto de datos HuggingFaceH4/ultrachat_200k, y que se ejecutó en 8 chips Trainium2, es decir, 32 NeuronCores, o sea, la mitad de una instancia trn2.48xlarge. El punto de guardado producido es un checkpoint Transformers estándar, utilizable en otros hardwares.
Varios pipelines de Diffusers son compatibles con TorchNeuron, tanto en modo « eager » como compilado. FLUX.1-dev, FLUX.2, Qwen-Image, Qwen-Image-2.1 y H3 también soportan el paralelismo de tensores, que reparte cada capa del modelo en varios núcleos. Hugging Face indica que otros pipelines serán progresivamente compatibles y que una guía permite añadir este soporte uno mismo.
Por último, la biblioteca Kernels acogerá núcleos de cálculo optimizados para Neuron, escritos con la Neuron Kernel Interface (NKI). Un núcleo destinado a los transformers de visión, publicado en el Hub, permite según la fuente una aceleración « hasta x2.4 » en Trainium. Hugging Face agradece al equipo de AWS por su trabajo en TorchNeuron y su reactividad ante los comentarios de la comunidad.
