AWS Trainium se vuelve tan fácil de usar como cualquier GPU con TorchNeuron
El nuevo plugin TorchNeuron elimina las barreras técnicas que hacían complicado entrenar modelos de IA en los chips Trainium de Amazon, permitiendo usar las mismas herramientas que con GPUs convencionales
TorchNeuron es un plugin de PyTorch (el framework de IA más popular) que permite usar chips Trainium de Amazon (procesadores diseñados para entrenar IA de forma más barata) con las mismas librerías y scripts que funcionan en GPUs convencionales, eliminando la necesidad de código especializado o herramientas extra como Optimum Neuron.
Amazon acaba de simplificar drásticamente el uso de sus chips de inteligencia artificial Trainium (procesadores diseñados específicamente para entrenar modelos de IA de forma más económica que con GPUs tradicionales). Según anunció Hugging Face (una plataforma donde desarrolladores comparten y descargan modelos de IA ya entrenados), el nuevo plugin TorchNeuron permite ejecutar modelos en estos chips usando exactamente las mismas librerías y scripts que se utilizarían en cualquier GPU de Nvidia o AMD, sin necesidad de código especializado.
Hasta ahora, trabajar con Trainium exigía pensar "como un compilador": había que usar PyTorch XLA (una versión especial de PyTorch, el framework de IA más popular), lidiar con tensores lazy (estructuras de datos que retrasan los cálculos hasta el último momento) en lugar de ejecución directa, y limitarse a formas estáticas que obligaban a recompilar el código cada vez que cambiaba algo. Hugging Face había creado Optimum Neuron, una capa extra de software, precisamente para ocultar esa complejidad y que los desarrolladores no tuvieran que aprender un sistema completamente distinto.
TorchNeuron cambia las reglas del juego al integrarse directamente en PyTorch y exponer Trainium como un dispositivo nativo más, igual que CUDA para GPUs. Esto significa que librerías como Transformers (para modelos de lenguaje), Diffusers (para generación de imágenes), PEFT (para ajuste fino eficiente) y TRL (para entrenar con refuerzo) funcionan directamente en Trainium sin modificaciones. En inferencia (cuando un modelo ya entrenado genera respuestas), el 89% de las 100 arquitecturas de modelos de texto más usadas ahora son compatibles de fábrica, frente al 46% que funcionaba con la versión anterior.
Para entrenar modelos, el cambio es igual de notable. Un script estándar de fine-tuning (el proceso de especializar un modelo preentrenado en una tarea concreta) escrito con TRL funciona en Trainium con apenas tres ajustes menores: desactivar las formas dinámicas automáticas de torch.compile (la herramienta de optimización de PyTorch), configurar el paralelismo de datos y tensores mediante DistributedConfig, y desactivar la carga asíncrona de pesos. El mismo código sirve para cualquier hardware, y el programador ya no necesita versiones especiales de clases o trainers (programas que gestionan el entrenamiento) dedicados a Neuron.
El plugin estará disponible como beta pública a partir del 20 de octubre de 2026 y se instalará simplemente con pip install torch-neuronx. Hugging Face recomienda que, de ahora en adelante, quienes usen Trainium o Inferentia (el chip de Amazon para inferencia rápida y barata) trabajen directamente con las librerías estándar en lugar de pasar por Optimum Neuron. La jugada de Amazon es clara: eliminar la fricción técnica que ha frenado la adopción de sus chips frente a las GPUs de Nvidia, apostando por que los desarrolladores elijan Trainium simplemente porque es más barato, no porque requiera un esfuerzo extra de aprendizaje.
