Hasta el 20 de octubre de 2026, llevar un modelo de PyTorch a una TPU obligaba a pasar por PyTorch/XLA o a reescribirlo en JAX. Ese día, en la PyTorch Conference 2026, Google publicó como open source TorchTPU, un backend nativo de PyTorch desarrollado con el equipo de PyTorch de Meta. La TPU pasa a ser un dispositivo más, como cuda o cpu, y Hugging Face ya lo integra en Transformers v5 y Diffusers.
Para quien entrena o sirve modelos en Google Cloud, el cambio importa más que el anuncio en sí: la pregunta "¿merece la pena portar a TPU?" deja de empezar por "¿cuánto código hay que tocar?".
Qué cambia frente a PyTorch/XLA
PyTorch/XLA (torch_xla) funcionaba, pero con un modelo mental propio. Hugging Face lo resume bien en su anuncio: los tensores eran perezosos, así que el código no siempre se ejecutaba cuando creías, cambios pequeños provocaban recompilaciones del grafo y depurar era difícil. Para salvar esa distancia, Hugging Face mantenía Optimum TPU, que con TorchTPU deja de ser necesario para ejecutar Transformers y Diffusers.
TorchTPU invierte el planteamiento:
| PyTorch/XLA | TorchTPU | |
|---|---|---|
| Modelo de ejecución | Tensores perezosos y grafo XLA | Modo eager por defecto, torch.compile para producción |
| Dispositivo en el código | API propia de torch_xla | "tpu", como cualquier otro dispositivo |
| Depuración | Recompilaciones difíciles de seguir | Eager mode y diagnósticos estándar de PyTorch |
| Hugging Face | A través de Optimum TPU | Nativo en Transformers v5 y Diffusers |
La idea es depurar en eager y pasar a torch.compile cuando toca exprimir la TPU. El backend trae kernels Pallas ajustados para TPU y primitivas distribuidas nativas.
El ecosistema que llega el mismo día
Lo que hace creíble el lanzamiento es que no llega solo. Google anunció a la vez integraciones en los motores que ya usa la mayoría:
- vLLM, con Inferact: backend nativo con batching continuo, cuantización y paralelismo multichip, sin cambiar workers, schedulers ni servidores de API.
- SGLang, con RadixArk, para cargas agénticas y estructuradas que reutilizan la caché KV entre peticiones.
- TorchTitan, para entrenamiento distribuido. Google cita hasta un 52 % de MFU (Model Flops Utilization) con Qwen 3 8B y escalado casi lineal hasta 8.192 TPU Ironwood.
- VeRL (en preview) para aprendizaje por refuerzo y Ray para orquestar CPU, GPU y TPU en un mismo flujo.
- llm-d 1.0 para enrutar inferencia en Kubernetes entre TPU y GPU con un endpoint común en GKE.
Hugging Face aporta su propio termómetro: de los tests de Transformers que se ejecutan en TPU pasa el 96,5 %, y de los 34 issues que han abierto en el repositorio de TorchTPU ya hay 19 corregidos. Es una cifra honesta: no todo está en verde, y lo dicen.
Cómo instalarlo y probarlo
En una VM con TPU, la instalación es un único pip. El índice extra sirve para traer la build de PyTorch para CPU; sin él, pip descarga la de CUDA, mucho más pesada e inútil en una TPU:
pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpuCon Transformers v5, cargar un modelo en la TPU es pasar device_map="tpu":
from transformers import AutoModelForCausalLM, AutoTokenizer
model_id = "Qwen/Qwen3.5-9B"
model = AutoModelForCausalLM.from_pretrained(model_id, device_map="tpu")
tokenizer = AutoTokenizer.from_pretrained(model_id)
inputs = tokenizer("The future of AI is", return_tensors="pt").to("tpu")
out = model.generate(**inputs, max_new_tokens=64)
print(tokenizer.decode(out[0], skip_special_tokens=True))En Diffusers basta con pipe.to("tpu") para el modo eager, y una llamada a compile(backend="tpu") sobre el transformer para el modo compilado. En los ejemplos de Hugging Face sobre una TPU v6e-4, una imagen de 1024×1024 con 40 pasos baja de 8,0 s en eager a 3,5 s compilada. La primera llamada compila; las siguientes con el mismo tamaño reutilizan el grafo.
Checklist antes de llevarlo a producción
Si ya tienes cargas en torch_xla u Optimum TPU, esto es lo que validaría antes de cambiar:
- Paridad numérica. Ejecuta el mismo lote en el stack actual y en TorchTPU, en eager, y compara salidas con una tolerancia razonable para
bfloat16. - Recompilaciones. En modo compilado, fija formas de entrada (
dynamic=Falseen los ejemplos de Diffusers) y mide cuántas veces compila en un tráfico real; la mejora de 8,0 a 3,5 s solo existe si el grafo se reutiliza. - Motor de serving. Si sirves con vLLM o SGLang, prueba su backend TorchTPU con tu modelo concreto antes de dar por hecho que hereda todas las funciones de la versión GPU.
- Issues abiertos. Revisa los issues del repositorio de TorchTPU que afecten a tu arquitectura: con un 96,5 % de tests pasando, el 3,5 % restante puede ser justo tu modelo.
Qué significa
Hasta ahora, elegir TPU era elegir también un stack (JAX o XLA). Con TorchTPU, la decisión vuelve a ser de precio y disponibilidad de hardware, que es como debería ser. La cita de Meta en el anuncio va en esa dirección: no quieren depender de una sola arquitectura. Para los equipos que solo trabajan con GPU de NVIDIA, la TPU entra en la lista corta sin coste de reescritura; lo que queda por demostrar es que el 52 % de MFU de TorchTitan se repite fuera de los benchmarks de Google.
Fuentes
- TorchTPU is now open source: Run native PyTorch on TPUs, Google Cloud Blog
- Hugging Face ❤️ TorchTPU: TPUs, Natively, Hugging Face