PyTorch ya funciona de forma nativa en los procesadores TPU de Google
TorchTPU elimina la principal barrera de software que impedía usar los aceleradores de inteligencia artificial de Google con código estándar de PyTorch, sin necesidad de aprender JAX ni lidiar con las complicaciones de torch_xla.
TorchTPU es un nuevo backend que integra las TPU (los chips de inteligencia artificial de Google) directamente en PyTorch, sin necesidad de usar torch_xla ni aprender JAX. Permite cargar y ejecutar modelos de Hugging Face en estos procesadores con un simple cambio de dispositivo, igual que se hace con GPUs.
Hasta ahora, si querías ejecutar código de PyTorch en las TPU (Tensor Processing Unit, los chips propios de Google diseñados específicamente para acelerar el cálculo matemático de las redes neuronales), tenías dos opciones poco cómodas: adaptar tu código a torch_xla (una capa intermedia con comportamientos poco intuitivos, como tensores que no ejecutaban operaciones cuando esperabas) o aprender JAX desde cero. TorchTPU, un nuevo backend nativo para PyTorch desarrollado por Hugging Face y Google, convierte las TPU en un dispositivo más dentro del ecosistema PyTorch. Ahora basta con escribir .to("tpu") para que un modelo se ejecute en estos procesadores, igual que harías con una GPU.
Las TPU llevan una década en desarrollo dentro de Google y se usan para entrenar y ejecutar modelos como Gemini. Están disponibles en Google Cloud desde un solo chip hasta configuraciones de miles de ellos, y destacan cuando necesitas mucha capacidad de cómputo, ya sea para entrenar modelos grandes o para servir predicciones a escala. El problema histórico ha sido el software: el ecosistema TPU creció alrededor de TensorFlow y JAX, dejando a los usuarios de PyTorch en una posición incómoda. Hugging Face desarrolló Optimum TPU para facilitar el trabajo, pero con TorchTPU esa solución intermedia ya no es necesaria.
La integración es completa: Transformers (la biblioteca de modelos preentrenados de Hugging Face) soporta TPU de forma nativa desde la versión 5, y Diffusers (la biblioteca para modelos de generación de imágenes y vídeo) también funciona directamente en estos chips. Cargar un modelo es tan simple como añadir device_map="tpu" al llamar a from_pretrained. Si el modelo es demasiado grande para un solo chip, puedes usar TensorParallelConfig para dividirlo automáticamente entre varios. Los primeros benchmarks muestran mejoras sustanciales: un kernel especializado de Tokamax para modelos de mezcla de expertos (MoE, una arquitectura donde distintas subredes se activan según la entrada) acelera la decodificación 21,5 veces en Qwen3-30B y 12,8 veces en Gemma 4 26B cuando se procesan lotes de 64 entradas.
Hugging Face está ejecutando su batería de tests de Transformers de forma continua en TPU, con un 96,5% de pruebas superadas. Los fallos detectados se reportan directamente al repositorio de TorchTPU (34 issues abiertos, 19 ya resueltos) y algunos han destapado bugs en el propio Transformers. La instalación en una máquina virtual con TPU requiere un solo comando: pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu. El índice extra sirve la versión CPU de PyTorch, mucho más ligera que la de CUDA, que no tiene sentido en este contexto. Según informó Hugging Face en su blog, el objetivo es que los desarrolladores descubran cualquier incompatibilidad antes de que afecte a proyectos reales.
