Hugging Face añade soporte nativo de PyTorch para los chips AWS Trainium

Hugging Face ha presentado TorchNeuron, un nuevo backend nativo que permite ejecutar modelos de PyTorch directamente en el hardware de AWS Trainium. Lanzado como beta pública el 20 de octubre de 2026, esta actualización elimina la necesidad de flujos de trabajo de compilación especializados, permitiendo a los desarrolladores utilizar las API estándar de PyTorch para entrenamiento e inferencia en los aceleradores de IA de Amazon.

Qué ha ocurrido

Durante varios años, ejecutar modelos de Hugging Face en los chips AWS Trainium e Inferentia requería el uso de Optimum Neuron. Esta biblioteca actuaba como puente entre PyTorch y el hardware, dependiendo del paquete torch-neuronx, construido sobre PyTorch XLA. Aunque funcional, este enfoque obligaba a los desarrolladores a adoptar una mentalidad centrada en el compilador. Debían lidiar con tensores diferidos (lazy tensors) en lugar de ejecución inmediata, limitarse a formas estáticas y gestionar recompilaciones frecuentes cada vez que cambiaba el grafo computacional. Optimum Neuron mantenía código personalizado, incluidas clases de modelado y entrenadores específicos, para ocultar esta complejidad a los usuarios.

TorchNeuron cambia esta arquitectura al integrarse directamente en PyTorch. Expone neuron como un dispositivo nativo, similar a cómo los desarrolladores interactúan con GPUs CUDA o dispositivos CPU. Este cambio habilita la ejecución inmediata (eager execution), lo que significa que las operaciones se realizan instantáneamente al ser llamadas, en lugar de quedar en cola para una compilación posterior. También admite APIs distribuidas nativas como Fully Sharded Data Parallel (FSDP) y Tensor Parallelism (TP), junto con las primitivas distribuidas propias de PyTorch como DTensor y FSDP2. Además, funciona perfectamente con torch.compile, permitiendo optimizaciones de rendimiento sin salir del flujo de trabajo estándar de PyTorch.

Como resultado, las principales bibliotecas de Hugging Face, incluyendo Transformers, Diffusers, Accelerate, PEFT y TRL, ahora se ejecutan directamente en Trainium. Los desarrolladores pueden usar las mismas clases y scripts que emplearían en cualquier otro acelerador. Hugging Face recomienda ahora a los usuarios migrar desde Optimum Neuron y utilizar estas bibliotecas directamente con el nuevo backend. La actualización busca simplificar la experiencia del desarrollador y facilitar el intercambio de modelos en el Hugging Face Hub para quienes utilizan la infraestructura de AWS.

Cómo funciona

El mecanismo central detrás de TorchNeuron es su integración como backend nativo de PyTorch. En lugar de traducir las operaciones de PyTorch a través de una capa de compilador externa como XLA, TorchNeuron permite que PyTorch reconozca neuron como un dispositivo de primera clase. Esto habilita la ejecución en modo inmediato, donde cada operación se ejecuta al instante en el hardware. Esto contrasta con el anterior modelo de evaluación diferida, que requería construir un grafo computacional completo antes de que pudiera comenzar la ejecución.

Al admitir primitivas distribuidas nativas, TorchNeuron aprovecha las herramientas integradas de PyTorch para escalar modelos en múltiples chips. Funciones como FSDP2 permiten un uso eficiente de la memoria durante el entrenamiento mediante la fragmentación (sharding) de los parámetros del modelo entre dispositivos. Para la inferencia, el backend admite paralelismo de tensor, que divide las capas del modelo entre múltiples aceleradores para agilizar el procesamiento. El sistema también fomenta el uso de cachés estáticos para tareas de generación con el fin de minimizar la sobrecarga de recompilación, asegurando que el rendimiento siga siendo alto incluso con la flexibilidad de la ejecución inmediata.

Detalles clave

• TorchNeuron está disponible como beta pública a partir del 20 de octubre de 2026.
• Los desarrolladores pueden instalar el paquete utilizando el comando pip install torch-neuronx.
• El nuevo backend admite ejecución inmediata, torch.compile y APIs distribuidas nativas como FSDP y TP.
• En Transformers, el 89% de las arquitecturas de modelos de generación de texto más utilizadas en el top 100 son compatibles de forma nativa, frente al 46% con Optimum Neuron v0.4.5.
• Los usuarios pueden mover modelos a dispositivos Neuron pasando el argumento device_map=neuron en su código.
• Hugging Face recomienda utilizar directamente bibliotecas como Transformers, Diffusers y otras, en lugar de pasar por Optimum Neuron.

Por qué importa

Esta actualización reduce significativamente la barrera de entrada para ingenieros que desean utilizar AWS Trainium. Anteriormente, la necesidad de pensar como un compilador y gestionar formas estáticas creaba una curva de aprendizaje pronunciada. Al alinear la experiencia de Trainium con los flujos de trabajo estándar de PyTorch, Hugging Face permite a los equipos reutilizar código y conocimientos existentes. Esto reduce el esfuerzo de ingeniería necesario para portar modelos al hardware de AWS y minimiza el riesgo de errores asociados con pasos de compilación personalizados.

El aumento en el número de arquitecturas de modelos soportadas también expande la utilidad de Trainium para cargas de trabajo de producción. Con el 89% de los principales modelos de generación de texto soportados nativamente, las organizaciones pueden desplegar una gama más amplia de modelos de última generación sin esperar integraciones personalizadas. Esta paridad con otros aceleradores significa que la elección de hardware puede estar impulsada más por el costo y la disponibilidad que por restricciones de compatibilidad de software.

Qué puedes hacer

• Instala el nuevo backend ejecutando pip install torch-neuronx en tu entorno.
• Actualiza tus scripts de inferencia para usar device_map=neuron al cargar modelos con Transformers.
• Reemplaza las dependencias de Optimum Neuron por llamadas directas a bibliotecas de Hugging Face como Transformers y Diffusers.
• Implementa cachés estáticos para tareas de generación de texto para evitar recompilaciones frecuentes y mejorar el rendimiento.
• Prueba tus configuraciones actuales de entrenamiento distribuido utilizando primitivas nativas de PyTorch como FSDP2 y DTensor en instancias Trainium.
• Consulta el blog de PyTorch para obtener detalles técnicos adicionales sobre cómo TorchNeuron se integra con el ecosistema más amplio de PyTorch.