Guía de TorchTPU para PyTorch en TPUs: instalación y migración en 2026

Migrar un proyecto de PyTorch de cuda a tpu con TorchTPU puede requerir muy pocos cambios, pero que el resultado funcione bien depende del acceso a la TPU, el paquete elegido, la compilación y las dimensiones de los tensores.

Qué ofrece realmente TorchTPU

TorchTPU es un backend nativo de PyTorch para las TPUs de Google. No es una API de modelos alojada ni sustituye a una TPU VM. El anuncio de PyTorch habla de ejecución eager, torch.compile, primitivas distribuidas y compatibilidad con aproximadamente 450 operaciones de ATen; el repositorio público incluye kernels de ATen, caché de compilación, ejemplos distribuidos e integración con torch.compile.

La migración parte de una idea muy familiar:

# Before
model.to("cuda")

# First TPU port
model.to("tpu")

Pero esto sirve como prueba de portabilidad, no como garantía de rendimiento. El anuncio de PyTorch señala que el rendimiento máximo en una TPU sigue dependiendo de las dimensiones del modelo, la fragmentación y el uso de operaciones adecuadas para este hardware.

Qué necesitas antes de instalarlo

TorchTPU solo resulta útil cuando tienes una TPU de Google aprovisionada. El repositorio recomienda ejecutar sus notebooks en una TPU VM y considera que una instancia v6e de Google Compute Engine es el entorno más sencillo para empezar. Un portátil sin TPU puede instalar los paquetes, pero no verificar la ejecución en una TPU.

El README público admite Python 3.11 o superior e identifica el paquete torch_tpu==0.2.0. Utiliza un entorno aislado para evitar que el experimento herede una compilación de CUDA u otras dependencias ajenas.

Instala torch_tpu sin arrastrar CUDA

Crea un entorno virtual y después instala la versión de PyTorch para CPU junto con el paquete fijado de TorchTPU:

python -m venv .venv
source .venv/bin/activate

pip install torch_tpu==0.2.0 \
  --extra-index-url https://download.pytorch.org/whl/cpu

Según el README, si omites el índice adicional puedes terminar instalando una wheel de PyTorch para CUDA, que no es la dependencia adecuada para una TPU VM.

Para una primera migración, utiliza la wheel publicada. Las compilaciones desde código fuente requieren Bazel o Bazelisk, además de las opciones documentadas por el repositorio: --config=wheel_common y, cuando sea necesario, --config=no_rbe.

Verifica el backend con un programa mínimo de PyTorch

Ejecuta este código en la TPU VM:

import torch

x = torch.randn(2, 2, device="tpu")
y = torch.randn(2, 2, device="tpu")
z = x @ y

print(z)
print(z.device)

Si obtienes un resultado correcto y el dispositivo aparece como tpu, habrás confirmado que la reserva de tensores y la multiplicación de matrices han llegado al backend TPU. Si el dispositivo no está disponible, revisa primero el runtime de la VM y la instalación de los paquetes antes de tocar el código del modelo.

TorchTPU admite ejecución eager para depurar y torch.compile para ejecutar código compilado. El anuncio de PyTorch indica que su modo DeferAndFuse puede ser un 50% o más rápido que el modo eager estricto en las pruebas del proyecto; considéralo un dato del proyecto, no un resultado universal.

Migra un workload real de PyTorch

Empieza con el cambio mínimo que funcione y añade la configuración distribuida específica de TPU solo después de validar el recorrido en un único dispositivo.

1. Sustituye la selección del dispositivo por torch.device("tpu") o utiliza torch.accelerator.current_accelerator() cuando sea compatible.
2. Mueve el modelo y los tensores de entrada al mismo dispositivo.
3. Ejecuta un batch corto antes de activar la compilación o la ejecución distribuida.
4. Mantén estables los tamaños de batch y las dimensiones de los tensores durante las mediciones; cambiar las dimensiones puede volver a activar la compilación.
5. Prueba el recorrido representativo del modelo y el patrón distribuido antes de comprometerte con una migración más amplia.

Un fragmento distribuido de CUDA suele utilizar NCCL. En cambio, el ejemplo publicado por TorchTPU usa el backend de grupos de procesos para TPU:

import torch
import torch.distributed as dist

device = torch.device("tpu")
dist.init_process_group(backend="tpu_dist")
model = model.to(device)

El comando exacto para lanzar la ejecución distribuida depende del entorno TPU y del workload. Cambiar la cadena del dispositivo no configura automáticamente la fragmentación, la ubicación de los procesos ni los pipelines de entrada.

La primera ejecución no tiene por qué parecerse a las siguientes. La caché de compilación documentada de TorchTPU puede reutilizar grafos repetidos, así que registra por separado el tiempo de calentamiento y el rendimiento sostenido.

En qué está preparado TorchTPU y qué sigue pendiente

Google documenta integraciones de TorchTPU con vLLM, SGLang, TorchTitan y las herramientas de Hugging Face. Su anuncio sobre inferencia describe rutas nativas para TPU en vLLM y SGLang, incluido el movimiento de la caché KV y distintas funciones de serving. El anuncio de Google Cloud recoge el resultado de 52% de MFU para Qwen 3 8B con TorchTitan y afirma que ya se utiliza en producción.

En la práctica, los límites están menos pulidos que los titulares de los anuncios:

Cuándo tiene sentido adoptar TorchTPU

Vale la pena probar TorchTPU si ya tienes acceso a una TPU de Google, ejecutas modelos de PyTorch y buscas una ruta común para entrenamiento o inferencia. El caso de uso inicial más sólido es el de un equipo dispuesto a perfilar y ajustar las dimensiones de los tensores, en lugar de esperar que un cambio de una sola línea ofrezca automáticamente una eficiencia equivalente a la de una GPU.

Es mejor esperar antes de migrar un workload crítico si necesitas una matriz de compatibilidad madura, un comportamiento predecible en el arranque en frío o una función de la que depende tu stack pero que TorchTPU todavía no expone. Haz la migración por etapas: un modelo, un segmento de TPU y un conjunto de evaluación fijo. Mide por separado el tiempo de instalación, la primera compilación, el rendimiento sostenido y la corrección de los resultados.

Preguntas frecuentes sobre TorchTPU

¿TorchTPU es lo mismo que PyTorch/XLA?

No. TorchTPU es un backend nativo de PyTorch más reciente, diseñado en torno a la ejecución eager y torch.compile; PyTorch/XLA es la integración anterior con TPUs. El código existente basado en XLA puede requerir un análisis de migración independiente.

¿Puedo instalar TorchTPU en un portátil sin TPU?

Puedes instalar los paquetes, pero no validar la ejecución en TPU sin una infraestructura TPU de Google compatible. La configuración práctica del repositorio oficial presupone una TPU VM.

¿Por qué el primer paso de TorchTPU se ejecutó tan despacio?

La primera ejecución puede compilar un grafo y llenar las cachés. Mide por separado las iteraciones posteriores y mantén estables las dimensiones mientras investigas posibles recompilaciones.

¿Cambiar cuda por tpu garantiza el mismo rendimiento?

No. Sustituir el dispositivo sirve para comprobar la corrección de la migración; antes de comparar el rendimiento debes medir las dimensiones, la fragmentación, la cobertura de operadores y el calentamiento.

¿TorchTPU está listo para producción?

Google habla de usos de inferencia y entrenamiento en producción, pero los componentes del ecosistema y sus funciones tienen distintos niveles de madurez. Prueba el modelo, el motor de serving y el patrón distribuido exactos antes de comprometerte.

Instala la wheel fijada en una TPU VM, ejecuta la verificación de 2×2 y compara un workload ya calentado con tu ruta actual en GPU. Así podrás separar la portabilidad del código, el calentamiento del compilador, la cobertura de operadores y la economía del hardware.