En el marco de la PyTorch Conference North America 2026, Google y AWS dieron un paso importante para diversificar las plataformas disponibles para los desarrolladores de inteligencia artificial basados en PyTorch. Google liberó públicamente TorchTPU, un stack nativo de PyTorch para sus TPUs, mientras que AWS lanzó TorchNeuron, un backend nativo para PyTorch compatible con sus chips Trainium, ambos disponibles en GitHub y con el objetivo de ampliar el ecosistema más allá del dominio de Nvidia y su software CUDA.
Durante años, Nvidia ha sido la opción predominante para el desarrollo de IA debido a la fortaleza de su ecosistema CUDA. Sin embargo, el lanzamiento de estas herramientas por parte de Google y AWS busca facilitar la migración a sus infraestructuras especializadas, aunque aún persisten desafíos para alcanzar un rendimiento óptimo sin realizar ajustes específicos de hardware.
Google apuesta por una integración más natural para PyTorch en TPUs
Google comenzó a trabajar en TorchTPU desde mediados de 2025, lanzando una versión preliminar en abril para algunos usuarios seleccionados. Esta iniciativa, desarrollada en conjunto con el equipo de PyTorch de Meta, pretende simplificar la ejecución de modelos PyTorch en sus chips TPU sin la necesidad de depender de JAX, reduciendo así el trabajo necesario para migrar cargas existentes a solo cambios mínimos en el código. 
Con TorchTPU, los desarrolladores pueden utilizar la ejecución eager para depurar modelos y torch.compile para optimizar el uso del TPU, disminuyendo el consumo de memoria. Esta integración viene acompañada de otras colaboraciones e integraciones anunciadas en la conferencia, entre las que destacan asociaciones con Hugging Face, vLLM, SGLang, llm-d, TorchTitan, Ray y TorchRec, así como aportes al ecosistema PyTorch como la contribución de LiteRT.
AWS también se suma con un backend nativo para Trainium
Por su parte, AWS presentó TorchNeuron, un backend nativo en PyTorch para sus chips Trainium, específicamente en versión beta pública que soporta las instancias Trainium2 y Trainium3. Al igual que TorchTPU, quienes deseen migrar modelos a Trainium podrían hacerlo con modificaciones muy sencillas, usando el método .to('neuron') en ejecución eager y asegurándose de seleccionar el backend Neuron en torch.compile cuando compilen el modelo.
TorchNeuron también se integra con Neuron Kernel Interface (NKI), un lenguaje y compilador de bajo nivel para programar dispositivos Neuron de AWS, lo que permite a los desarrolladores crear operadores personalizados y optimizar el rendimiento específico en Trainium. AWS junto con PyTorch han publicado además guías detalladas que explican los modos de ejecución y los compromisos de rendimiento de esta nueva herramienta.
Un avance significativo pero con margen de mejora
Aunque estas soluciones representan un avance crucial para ofrecer más alternativas a la tradicional dependencia en Nvidia, por el momento no ofrecen una migración completa ni un rendimiento máximo sin optimizaciones específicas para cada plataforma. Según PyTorch, aunque una modificación mínima en el código puede ser suficiente para que un modelo funcione en TPUs, alcanzar el máximo rendimiento exige optimizaciones puntuales y el uso de stacks de producción especializados.
En particular, para tareas distribuidas, el mecanismo predeterminado Skinny Collectives en TorchTPU puede ocasionar una pérdida de rendimiento de entre un 30% y un 40% debido a que limita la capacidad del compilador XLA para solapar computación y comunicación, por lo que se recomiendan alternativas más optimizadas como Pallas kernels o stacks de producción como vLLM-TorchTPU, SGLang-TorchTPU y TorchTitan-TorchTPU.
En cuanto a AWS y Trainium, TorchNeuron se encuentra en beta pública y, aunque en la mayoría de casos migrar modelos requiere solo un cambio de configuración, algunos proyectos de alto rendimiento, como el caso de i14 adaptando su modelo Laplace 3B, demandaron ajustes más profundos en el código para aprovechar la eficiencia de Trainium2.
En síntesis, Google y AWS han abierto nuevas sendas para que desarrolladores de PyTorch puedan explorar alternativas tecnológicas a Nvidia sin comenzar desde cero, aunque el camino para explotar completamente el potencial de estas plataformas sigue requiriendo trabajo y especialización.
