Google anunció TorchTPU el 20 de octubre de 2026 como backend de código abierto para ejecutar PyTorch en sus TPU. Para una ejecución básica, el cambio puede reducirse a instalar el paquete y seleccionar el dispositivo TPU; los trabajos distribuidos requieren además un backend y una configuración de lanzamiento específicos.
El cambio básico: instalar y seleccionar la TPU
La documentación de Google Cloud indica instalar TorchTPU con pip install torch_tpu y elegir torch.device("tpu"). El equipo de PyTorch ilustra el cambio de .to("cuda") a .to("tpu"). La ejecución eager, que ejecuta las operaciones directamente, es la opción predeterminada; torch.compile permite compilar el modelo.
TorchTPU, desarrollado por Google junto con el equipo de PyTorch de Meta, busca que el código PyTorch pueda ejecutarse en TPU con cambios acotados. Cambiar el dispositivo puede bastar para empezar, pero no garantiza el máximo rendimiento: las formas de los tensores, la distribución del trabajo entre chips y las operaciones pueden estar optimizadas para GPU.
| Tipo de ejecución | Dispositivo o backend | Lanzamiento |
| Básica | torch.device("tpu") o .to("tpu") | Instalación de torch_tpu; ejecución eager predeterminada o compilación con torch.compile |
| Distribuida | Backend tpu_dist | Un proceso por dispositivo y configuración de topología TPU y slice-builder |
Los trabajos distribuidos necesitan configuración TPU
En código distribuido, TorchTPU usa el backend tpu_dist. El modelo previsto es un proceso por dispositivo, y el lanzamiento requiere configurar la topología TPU y los valores de slice-builder, o utilizar el helper documentado por el proyecto. Por eso, un cambio de dispositivo no basta para trasladar por sí solo un trabajo distribuido.
También puede haber que adaptar operaciones concretas. En un ejemplo de portabilidad de Muse Glimmer 30B, el equipo de PyTorch sustituyó FlexAttention, que no estaba soportada por TorchTPU en esa prueba, por F.scaled_dot_product_attention.
Qué indican las pruebas publicadas por PyTorch
El proyecto informa que TorchTPU admite alrededor de 450 operadores ATen, sin contar delegados estructurales o compuestos. En su prueba Tier 1, pasaron la comprobación de ejecución sin caídas 4.821 de 4.917 modelos seleccionados, un 98 % según el proyecto. En una muestra distinta, Tier 2, 638 de 657 arquitecturas representativas alcanzaron paridad numérica frente a una referencia GPU, un 97 % informado.
En su evaluación de dinamismo acotado, el proyecto obtuvo resultados correctos en 834 de 839 modelos de LLM causales, un 99 %. Son resultados de esos conjuntos de prueba; no constituyen una garantía para cualquier modelo ni una comparación estandarizada de rendimiento entre GPU y TPU.
El equipo también informó de un entrenamiento de Qwen3 MoE 235B escalado hasta 8.192 chips TPU v7. En esa configuración, con una longitud de secuencia de 4.096 y un lote de 262.144, comunicó 2.104 tokens por segundo por TensorCore, un 27,1 % de MFU y 31,2 segundos por paso. Es una medición del proyecto para ese modelo y esa configuración.
Cargar pesos guardados desde una GPU
Es posible cargar en TPU los pesos de un modelo PyTorch compatible mediante un state_dict y reasignar el dispositivo con torch.load. La documentación de PyTorch sobre guardar y cargar modelos explica cómo restaurar esos parámetros; la documentación de TorchTPU describe el uso de semántica PyTorch estándar en la TPU.
Los pesos del modelo son solo una parte de un checkpoint de entrenamiento. El estado del optimizador y del planificador se guarda aparte, así que cargar el state_dict del modelo no restaura automáticamente todo el estado necesario para reanudar el entrenamiento.
