Los desarrolladores que trabajan con PyTorch ya pueden llevar sus modelos a las TPU de Google sin pasar por JAX ni por capas de compatibilidad: Google ha liberado TorchTPU, un backend nativo de PyTorch para sus aceleradores, con licencia Apache 2.0. La compañía lo presentó el 20/10/2026 en la PyTorch Conference 2026 y lo ha desarrollado junto al equipo de PyTorch de Meta. Llega con soporte en Hugging Face Transformers, vLLM y SGLang desde el primer día.
Las claves de TorchTPU para desarrolladores de PyTorch en 30 segundos
- TorchTPU es código abierto (Apache 2.0) y la versión 0.2.0 ya se instala con pip.
- Funciona como un dispositivo más de PyTorch, con depuración en modo eager y compilación con torch.compile.
- Hugging Face Transformers v5 permite cargar modelos en TPU con device_map=»tpu».
- vLLM y SGLang tienen backends nativos para servir modelos en TPU.
- Google asegura un 52 % de MFU entrenando Qwen 3 8B en TPU Ironwood, según sus pruebas.
Hasta ahora, usar PyTorch en una TPU significaba recurrir a PyTorch/XLA, al que Google presenta ahora como predecesor de TorchTPU. El nuevo proyecto se conecta a PyTorch mediante PrivateUse1, la interfaz que el framework ofrece para backends de hardware de terceros. Para el desarrollador, eso se traduce en tensores normales de PyTorch que viven en la TPU, sin clases especiales.
Cómo se usa TorchTPU
El paquete se instala desde PyPI y requiere Python 3.11 o superior:
pip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpuGoogle propone dos fases. Primero, depurar en modo eager, con tres variantes según el blog técnico del equipo: una de depuración que sincroniza con la CPU tras cada operación (lenta, pero útil para encontrar NaN, errores de forma o falta de memoria), otra estricta que imita el comportamiento normal de PyTorch y una que fusiona operaciones automáticamente, con una mejora de entre el 50 % y más del 100 % sobre la estricta, según Google. Después, para producción, torch.compile captura el grafo y lo compila con XLA, el compilador de Google para TPU.
Para entrenamiento distribuido, TorchTPU admite DDP, FSDPv2 y DTensor, y permite escribir kernels a medida en Pallas o JAX. El repositorio incluye cuadernos de ejemplo pensados para una instancia con TPU v6e.
Hugging Face, vLLM y SGLang, integrados de serie
La parte que más notarán muchos usuarios está en las bibliotecas que ya usan. En Transformers v5, Google explica que desplegar en TPU es tan sencillo como pasar device_map="tpu" al cargar el modelo, y Diffusers funciona de forma nativa con modo eager, torch.compile y paralelismo de tensores. Hugging Face ejecuta cada noche su batería de pruebas de Transformers sobre TPU. Jeff Boudier, vicepresidente de producto de Hugging Face, defiende en el anuncio que la diversidad de hardware es esencial para que la IA de código abierto sea accesible para todos.
Para servir modelos, vLLM tiene un backend nativo sobre TorchTPU desarrollado con Inferact, con kernels Pallas propios, batching continuo, cuantización y paralelismo entre varios chips. SGLang, con RadixArk, aporta su RadixAttention para reutilizar la caché KV entre peticiones. Y en entrenamiento, TorchTitan ofrece recetas para Llama 3, Qwen 3 y FLUX.1, mientras que VeRL incorpora una receta de aprendizaje por refuerzo con GRPO, en fase preliminar. Google también ha publicado en vista previa un conjunto de complementos para asistentes de programación que ayudan con la puesta en marcha, el diagnóstico y el ajuste de rendimiento en TPU.
Lo que dice Google del rendimiento y lo que falta
Google afirma haber alcanzado un 52 % de utilización de FLOPs del modelo (MFU) entrenando Qwen 3 8B con TorchTitan, con un escalado casi lineal hasta 8.192 TPU Ironwood. Son datos de la propia empresa, sin pruebas independientes todavía. Meta, socio del proyecto, asegura a través de Jana van Greunen, directora de ingeniería de PyTorch, que TorchTPU le permite llevar sus cargas de PyTorch a las TPU de Google, incluidos modelos de lenguaje y sistemas de anuncios y recomendación.
El propio plan de trabajo del equipo reconoce los puntos débiles: reducir las recompilaciones cuando cambian el tamaño del lote o la longitud de las secuencias, dar soporte completo a formas dinámicas en torch.compile y crear una biblioteca de kernels precompilados para acortar la primera ejecución. Y hay un límite evidente: TorchTPU solo sirve para las TPU de Google, que se usan a través de Google Cloud. En noticias.ai ya se explicó qué chips compiten por ejecutar la IA; con TorchTPU, Google intenta que elegir TPU no obligue a cambiar de framework.
Preguntas frecuentes
¿Qué es TorchTPU?
Un backend nativo de PyTorch para las TPU de Google, publicado como código abierto con licencia Apache 2.0 y desarrollado con el equipo de PyTorch de Meta.
¿Cómo se instala TorchTPU?
Con pip: pip install torch_tpu==0.2.0 y el índice de PyTorch para CPU como índice adicional. Requiere Python 3.11 o superior.
¿Funciona Hugging Face Transformers con TorchTPU?
Sí. En Transformers v5 basta con indicar device_map=»tpu» al cargar el modelo, según Google. Diffusers también funciona de forma nativa.
¿Hace falta reescribir el código de PyTorch?
Google asegura que los cambios de código son mínimos y que se usan las API habituales de PyTorch, con modo eager para depurar y torch.compile para producción.
Fuentes:
