Google open sourcea TorchTPU, un backend nativo de PyTorch para Cloud TPUs

Google ha open sourceado TorchTPU, un backend nativo de PyTorch para Cloud TPUs. Fue codesarrollado con el equipo de PyTorch de Meta y anunciado en la PyTorch Conference 2026 el 20 de octubre de 2026. El código PyTorch existente apunta a una TPU a través de un estándar torch.device("tpu"). Las TPUs ahora se sitúan junto a las GPUs y CPUs como un dispositivo nativo de PyTorch. El equipo de Google declara que TorchTPU reemplaza a PyTorch/XLA y lo recomienda para todas las nuevas cargas de trabajo.

TL;DR

• Tamaño: Una pila de software, no un modelo. 4 repositorios abiertos: TorchTPU, vLLM-TorchTPU, SGLang-TorchTPU, TorchTitan. TorchTPU es v0.2.0 en PyPI, Python 3.11+, Apache 2.0.
• Se ejecuta en: Google Cloud TPUs. Probado en generaciones desde v5e, escalado hasta 8,192 chips Ironwood (TPU7x). El serving admite checkpoints FP8, MXFP4 y NVFP4.
• Rendimiento: Modo eager para depuración, torch.compile para rendimiento a nivel de producción.
• Mejor: 52% MFU entrenando Qwen 3 8B con TorchTitan en Ironwood (autoinformado).
• Peor: Eager sin ajustes tomó 11.49s por paso en un caso de fine-tuning, frente a 2.78s completamente ajustado.

¿Qué es TorchTPU?

TorchTPU es un backend de PyTorch open source de Google y Meta que ejecuta modelos PyTorch de forma nativa en TPUs. Se conecta a través de la interfaz PrivateUse1de PyTorch. Los tensores son tensores ordinarios de PyTorch, sin subclases ni wrappers. Ejecuta un proceso por dispositivo y preserva las vistas de tensores y las mutaciones in-place.

El predecesor, PyTorch/XLA, registraba grafos de tensores lazy. Los desarrolladores a menudo añadían código para marcar regiones compilables. PyTorch/XLA también solo admitía código SPMD puro. TorchTPU también maneja ejecución MPMD divergente.

¿Cómo ejecuta TorchTPU el código de PyTorch?

TorchTPU incluye 3 modos eager:

• Strict (predeterminado): despacha cada operador de forma asíncrona con la semántica estándar de PyTorch.
• Debug: se ejecuta de forma síncrona con comprobación de límites de tensores.
• DeferAndFuse: fusiona rachas de operadores eager para reducir el tráfico de HBM. El artículo de TorchTPU indica que puede ser un 50% o más rápido que Strict.

Para máxima velocidad, torch.compile(backend="tpu") captura el grafo con TorchDynamo. Lo reduce mediante StableHLO a OpenXLA, que emite código máquina de TPU. Los ejecutables compilados se almacenan en caché entre procesos, hosts y reinicios de trabajos.

Los kernels personalizados de Pallas se pueden invocar como funciones estándar de PyTorch. Los kernels preconstruidos de Tokamax y los kernels de tiles de Helion también funcionan.

La distribución usa el tpu_dist grupo de procesos. Admite DDP, FSDP2 y DTensor, junto con una mesh consciente de la topología para las interconexiones de TPU.

¿Qué tan rápido es TorchTPU?

Todas las cifras siguientes son autorreportadas por Google y sus socios.

• Escala de entrenamiento: TorchTitan alcanzó hasta un 52% de MFU en Qwen 3 8B. Google informa un escalado casi lineal hasta 8,192 TPU Ironwood.
• Escalera de ajuste: En un caso de fine-tuning con TorchTitan, el tiempo por paso bajó de 11.49s a 2.78s. DeferAndFuse por sí solo dio 1.15x. La configuración completa torch.compile dio 3.80x.
• Cobertura de modelos: Google prueba continuamente alrededor de 5,000 modelos de Hugging Face en modo eager y compilado.
• Sirviendo en flota: Zenteiq aumentó su producción diaria de 1.47B a 5.97B tokens por cada 1,000 chips TPU v6e. Esa ganancia provino del enrutamiento de llm-d en GKE, no de TorchTPU por sí solo.

¿Qué integraciones se incluyen con TorchTPU?

• Hugging Face: Transformers v5 se despliega en TPU con device_map="tpu". Diffusers soporta eager, torch.compile y paralelismo de tensores.
• vLLM-TorchTPU: construido con Inferact. Mantiene los planificadores upstream y los servidores de API, con kernels Pallas personalizados.
• SGLang-TorchTPU: construido con RadixArk. Mantiene la reutilización de prefijos de RadixAttention.
• Características de serving: Cachés KV en FP8, checkpoints FP8, MXFP4 y NVFP4. La decodificación especulativa incluye MTP y EAGLE-3.
• TorchTitan: recetas para Llama 3 70B/405B, Qwen 3 8B/32B, Qwen 3 235B MoE y FLUX.1 12B.
• Además: llm-d 1.0, Ray, TorchRec, VeRL (vista previa) y LiteRT-Torch para exportación en el dispositivo.

¿Cómo se compara TorchTPU con PyTorch/XLA, TorchNeuron y JAX?

¿Cómo se instala TorchTPU?

Instala el paquete en una TPU VM y luego selecciona el tpu dispositivo:

# pip install torch_tpu
import torch
device = torch.device("tpu")
x = torch.randn((4, 4), device=device)

El README del repositorio fija torch_tpu==0.2.0 con el índice de wheels de PyTorch para CPU. Sin ello, pip descarga la build CUDA de torch. Las guías completas están en torchtpu.readthedocs.io.

Conclusiones clave

• TorchTPU es de código abierto bajo Apache 2.0, construido junto con Meta.
• Reemplaza PyTorch/XLA con una ejecución eager-first basada en PrivateUse1.
• torch.compile(backend="tpu") reduce el paso por StableHLO y XLA.
• TorchTitan alcanzó un 52% de MFU en Qwen 3 8B, escalando hasta 8,192 chips Ironwood.
• vLLM, SGLang y Hugging Face ahora se ejecutan en TPU a través de PyTorch estándar.

Consulta el repositorio de GitHub, artículo, documentación y anuncio de Google Cloud. Todo el crédito corresponde al investigador de este proyecto. Además, no dudes en seguirnos en Twitter y no olvides unirte a nuestro ML SubReddit de más de 150k y suscribirte a nuestro boletín. ¡Espera! ¿estás en telegram? ahora también puedes unirte a nosotros en telegram.

La publicación Google hace open-source de TorchTPU, un backend nativo de PyTorch para Cloud TPUs apareció primero en MarkTechPost.