TorchTPU: Google libera como código abierto un PyTorch nativo para TPU

Google liberó TorchTPU como código abierto el 20 de octubre de 2026, un backend que ejecuta código de PyTorch de forma nativa en sus chips TPU. Importa porque sustituye a PyTorch/XLA, la vía más antigua y difícil para llegar a las TPU, que se archivará el 15 de noviembre. Google afirma que el código existente debería migrar con poco más que un cambio del nombre del dispositivo.

Qué es TorchTPU

Las TPU son los chips de IA que Google diseña y alquila a través de Google Cloud. PyTorch es la biblioteca más usada para crear y entrenar modelos de IA. Hasta ahora, PyTorch llegaba a las TPU a través de PyTorch/XLA, que convierte un modelo en un grafo compilado antes de ejecutarlo.

En cambio, TorchTPU convierte la TPU en un dispositivo normal de PyTorch. «Portar una carga de trabajo a TPU debería ser tan sencillo como cambiar model.to("cuda") por model.to("tpu")», escribe el equipo en el anuncio del blog de PyTorch. Google lo desarrolló con el equipo de PyTorch de Meta, a partir de mediados de 2025.

El repositorio de GitHub del proyecto indica la versión 0.2.0 en PyPI, la licencia Apache 2.0 y Python 3.11 o posterior. Crypto Briefing informa de que el lanzamiento abarca cuatro repositorios y se presentó en la PyTorch Conference. El blog los nombra como TorchTPU, vLLM-TorchTPU, SGLang-TorchTPU y TorchTitan, que cubren el backend principal, dos motores para servir modelos y una biblioteca de entrenamiento.

Cómo ejecuta el código

TorchTPU implementa unas 450 operaciones ATen, las funciones matemáticas básicas con las que se construye todo modelo de PyTorch. Por defecto, el código se ejecuta línea a línea en modo «eager», como en PyTorch normal. El blog describe tres modos eager:

Modo	Cómo activarlo	Qué hace
Strict	Por defecto	Ejecuta cada operación tal como está escrita
Debug	TORCH_TPU_LAUNCH_BLOCKING=1	Espera a cada operación, para facilitar la depuración
DeferAndFuse	TORCH_TPU_DEFER_AND_FUSE=1	Agrupa las operaciones; «un 50 % o más rápido que Strict Eager»

Para ganar velocidad, torch.compile(backend="tpu") convierte un modelo en un único programa de TPU. El código personalizado para el chip puede escribirse con Pallas, de JAX, o con Helion. El blog enumera integraciones con vLLM, SGLang, TorchTitan, Ray y TorchRec. El blog también afirma que algunos de los mayores clientes de Google Cloud ya usan TorchTPU en producción.

Qué tal funcionan los modelos

Por ahora, las únicas cifras son las de las pruebas de la propia Google. El blog informa de estos resultados:

Prueba	Resultado
Modelos de Hugging Face ejecutados	4.917, de más de 500 familias de modelos
Superaron el primer nivel de pruebas	98 % (4.821 de 4.917)
La salida coincidió con la de una GPU	97 % (638 de 657 comprobados)
Entrenamiento de Qwen3 235B en 64 chips TPU v7	1.828 tokens/s en PyTorch, frente a 1.937 en la pila JAX de Google

Crypto Briefing cita la cifra de Google de un 52 % de utilización de FLOPs del modelo al entrenar Qwen 3 8B en hasta 8.192 chips TPU v7x. Esa medida es la parte de la capacidad de cálculo máxima de un chip que el modelo usa realmente.

PyTorch/XLA se retira

El repositorio de PyTorch/XLA ahora se abre con un aviso: «Este proyecto está obsoleto y ha sido sustituido por TorchTPU». Añade que el repositorio «se archivará el 15 de noviembre de 2026», lo que lo deja en modo de solo lectura. Su consejo es «empezar las nuevas cargas de trabajo de PyTorch en TPU con TorchTPU y migrar las cargas de trabajo existentes de PyTorch/XLA». La última instalación estable que indica fija torch y torch_xla en la versión 2.8.0.

El mismo día, el backend TorchNeuron de Amazon para sus chips Trainium entró en beta pública y abandonó el mismo enfoque XLA. Ambos proveedores de nube tratan ahora sus chips de IA como dispositivos corrientes de PyTorch.

Qué significa para los desarrolladores

TorchTPU convierte una TPU en una cadena de dispositivo más, lo que sitúa los chips de Google en la misma comparación de precios que las GPU.

• Planifica ya tu salida de PyTorch/XLA. Después del 15 de noviembre, el repositorio no acepta correcciones. Fija torch_xla 2.8.0 para los trabajos en curso y prueba TorchTPU en una máquina v6e, que el README sugiere como el punto de partida más fácil.
• Revisa primero el 2 %. La propia prueba de Google aún falla en aproximadamente uno de cada cincuenta modelos. Ejecuta un trabajo corto de entrenamiento o inferencia con tu modelo exacto antes de planificar un cambio.
• Reserva tiempo para los kernels de CUDA personalizados. El código de GPU escrito a mano no se traslada. Hay que reescribirlo en Pallas o Helion, así que incluye ese trabajo en el presupuesto.
• Fija la versión. Esta es la versión 0.2.0. Fíjala en los archivos de requisitos y vuelve a comprobar los resultados tras cada actualización.
• Haz tus propias pruebas de rendimiento. Todas las cifras de velocidad de este artículo proceden de Google. Compara el coste por token en tu propia carga de trabajo, frente a las GPU y los chips de otras nubes.