PyTorch en TPU y Trainium: ¿cuánto reduce la dependencia de Nvidia?

Ejecutar PyTorch en TPU y Trainium ya solo exige cambiar el dispositivo con TorchTPU y TorchNeuron. Analizamos cifras, límites, costes reales y quién debe probar.

Noticia

Ejecutar PyTorch en las TPU de Google o en los chips Trainium de AWS exigía hasta ahora una biblioteca intermedia y bastantes ajustes. El 19 y 20 de octubre de 2026, justo antes de la PyTorch Conference, ambas empresas liberaron un backend nativo de PyTorch: TorchTPU (Google, desarrollado con el equipo de PyTorch de Meta) y TorchNeuron (AWS, en beta pública). El objetivo común: que el código PyTorch escrito para GPU de Nvidia funcione en sus chips con muy pocos cambios.

• Cambiar el dispositivo, no reescribir. En TorchNeuron, AWS afirma que el único cambio es model.to('neuron'). En TorchTPU, Google habla de "cambios mínimos de código"; en Hugging Face Transformers v5 basta con device_map="tpu".
• Los dos modos de ejecución. Ambos admiten modo eager (operación a operación, fácil de depurar) y torch.compile (compilación del grafo completo para producción).
• Entrenamiento distribuido incluido. DDP, FSDP2 y DTensor funcionan en ambos; TorchNeuron añade paralelismo de tensores y de pipeline.
• Integrados en el ecosistema de inferencia. TorchTPU trae backends para vLLM y SGLang; TorchNeuron cuenta con vLLM Neuron y soporte en Transformers v5, Diffusers y trl.
• Código abierto Apache 2.0 para ambos proyectos en GitHub.

Análisis

Análisis: ¿basta una línea de código para dejar las GPU de Nvidia?

Respuesta corta: basta para probar, no para migrar. El cambio real está en la forma de integración; el ahorro depende mucho de los modelos que ejecute.

¿Qué cambia respecto a antes?

La diferencia clave es PrivateUse1, el punto de extensión oficial de PyTorch para hardware nuevo. Un tensor en TPU o Trainium es un tensor PyTorch normal, sin envoltorios. Por eso bibliotecas como Hugging Face, vLLM o TorchTitan ya no necesitan una rama distinta para cada chip.

¿Por qué funciona sin reescribir?

Al llamar a torch.compile, PyTorch usa TorchDynamo para capturar el grafo de cálculo. En lugar de pasarlo a un compilador de GPU, TorchTPU lo traduce a StableHLO y deja que XLA lo optimice para TPU. TorchNeuron sigue un camino parecido con Torch-MLIR, StableHLO y el compilador Neuron. El desarrollador sigue escribiendo PyTorch; la traducción al chip ocurre por debajo.

El modo "Fused Eager" de TorchTPU agrupa automáticamente muchas operaciones pequeñas en bloques mayores antes de enviarlas al chip. Google ha informado de mejoras del 50% a más del 100% frente a la ejecución operación a operación. Es una cifra del fabricante, sin medición independiente.

¿A quién afecta y cómo?

• Ingenieros de ML: pueden ejecutar el mismo bucle de entrenamiento en GPU, TPU y Trainium para comparar. Los kernels propios siguen exigiendo aprender Pallas (TPU) o NKI (Trainium).
• Responsables técnicos: ganan margen al negociar precios de GPU. Un contrato en la nube deja de estar atado a CUDA si los modelos corren en varios chips.
• Pymes: la mayoría usa modelos vía API, así que el efecto es indirecto. Los equipos que alojan modelos abiertos (Qwen, Llama) en Google Cloud o AWS se benefician directamente.
• Marketing y producto: nada cambia de inmediato. El beneficio, si llega, será vía precios de API más bajos de proveedores que usen estos chips.

¿Cuánto cuesta en la práctica?

Ninguna empresa publicó precios nuevos. La única cifra de coste es la de Deepgram: alrededor de un 60% más de audio procesado por dólar que con sus GPU actuales. La tabla es un ejemplo hipotético que aplica esa cifra a dos escalas, con un coste de migración supuesto de 30.000 USD (unos dos ingenieros durante 1–2 meses).

Conclusión: con una factura pequeña, cambiar de chip rara vez compensa. Con una factura grande y un modelo estable, las cuentas empiezan a salir. El 60% procede de un solo cliente y una sola carga (audio); su modelo puede rendir mejor, peor o no ganar nada.

Riesgos y límites

• Cobertura de operadores incompleta. En la batería de pruebas de PyTorch de agosto de 2026 se omitieron unas 22.000 pruebas porque TorchNeuron no tenía el operador correspondiente. La primera versión no admite formas dinámicas, tensores dispersos ni NestedTensor; FP8 está en beta privada.
• Las formas dinámicas siguen siendo un punto débil en TPU. Google reconoce que longitudes de secuencia o tamaños de lote variables provocan recompilaciones y que la primera iteración tiene alta latencia. Un chatbot con prompts de longitud muy variable es justo el tipo de carga afectada.
• La dependencia solo cambia de lugar. Las TPU solo existen en Google Cloud y Trainium solo en AWS. Salir de CUDA puede significar atarse más a una nube.
• El código exclusivo de CUDA no funciona. CUDA graphs, llamadas específicas de NCCL y kernels CUDA propios hay que reescribirlos.

Puntos por verificar

• El 52% de MFU y los 8.192 chips son pruebas internas de Google con TorchTitan, sin comparación independiente con GPU de precio similar.
• Las cifras de Zenteiq (unas 4 veces) y Deepgram (+60% por dólar) son testimonios de clientes elegidos por los fabricantes.
• El dato de "85% de las apps de IA usan PyTorch" citado por TechCrunch no tiene fuente identificada; TBN Grow no pudo verificarlo.
• Cómo comprobarlo usted mismo: tome un modelo en producción, mida tokens/segundo y coste por millón de tokens en sus GPU actuales y ejecute el mismo modelo en TPU v6e o Trainium2 con las mismas longitudes de entrada.

La valoración de TBN Grow

Es un cambio en el software de infraestructura con valor real para equipos que entrenan o alojan modelos abiertos en Google Cloud o AWS. Quien tenga facturas de GPU de decenas de miles de dólares al mes debería probarlo ya con un modelo estable. Equipos pequeños, usuarios de API o cargas con entradas muy variables deberían esperar una o dos versiones, hasta que maduren la cobertura de operadores y las formas dinámicas.

Qué observar

Qué vigilar al ejecutar PyTorch en TPU y Trainium

Lo que puede hacer ya

• [ ] Localice cada llamada directa a cuda en su código y sustitúyala por la API torch.accelerator para facilitar el cambio de chip.
• [ ] Pruebe TorchTPU: pip install torch_tpu==0.2.0 en una VM TPU v6e (el README del proyecto la señala como el inicio más sencillo; añada --extra-index-url https://download.pytorch.org/whl/cpu para no instalar la versión CUDA).
• [ ] Pruebe TorchNeuron: pip install torch-neuronx en una instancia trn2 y cambie el dispositivo a neuron.
• [ ] Ejecute las pruebas de su propio proyecto para detectar operadores no soportados antes de calcular costes.
• [ ] Compare el coste por millón de tokens o por hora de audio, no solo la velocidad.

Calendario y lo que aún no se ha anunciado

• AWS prevé llevar TorchAO y FP8 a disponibilidad general en la próxima versión y cerrar huecos de operadores; sin fechas.
• Google quiere añadir formas dinámicas a torch.compile y una biblioteca de kernels precompilados; sin calendario.
• Google no ha etiquetado TorchTPU como beta ni como versión general, ni ha publicado la lista completa de generaciones de TPU compatibles.
• Nvidia no ha respondido oficialmente.

Más lectura: El mercado chino de chips de IA cambia: Nvidia baja al 34% y Huawei sube al 28% y el problema de rentabilidad de la inversión en infraestructura de IA.

Fuentes de verificación: Google Cloud Blog, Google Developers Blog, TorchTPU en GitHub, PyTorch Blog – TorchNeuron, AWS What's New, TechCrunch.