Si alguna vez intentaste ejecutar PyTorch en una TPU, probablemente portaste tu código a torch_xla, o terminaste teniendo que aprender JAX. Con TorchTPU, un backend nativo de PyTorch, la TPU ahora es solo otro dispositivo. Esto abre el soporte nativo de TPU en el ecosistema de Hugging Face. Transformers, Diffusers y Kernels se ejecutan de forma nativa en TPU. Elige un modelo del Hub, envíalo a "tpu", y listo.
Espera, ¿qué es una TPU? 🤔
TPU significa Tensor Processing Unit. Las TPUs son los aceleradores de IA de Google. Son chips construidos para un solo trabajo: las operaciones matriciales en las que las redes neuronales pasan la mayor parte de su tiempo. Google las viene diseñando desde hace unos diez años, y las usa para entrenar y servir sus propios modelos, incluido Gemini.
Las TPUs están disponibles en Google Cloud, desde un solo chip hasta pods de miles de chips. Brillan cuando necesitas mucha potencia de cómputo, ya sea para entrenamiento o para servir a gran escala.
El inconveniente siempre ha sido el software. El stack de TPU creció en torno a TensorFlow y JAX, y los usuarios de PyTorch tenían que esforzarse más para acceder.
La manera antigua
Hasta ahora, ejecutar PyTorch en TPUs significaba pasar por PyTorch/XLA (torch_xla). Funcionaba, pero tenía sus peculiaridades. Los tensores eran perezosos (los tensores, no tú), así que tu código no siempre se ejecutaba cuando pensabas que lo hacía. Pequeños cambios podían desencadenar recompilaciones del grafo, y depurar era difícil de razonar.
Para facilitar la vida, creamos Optimum TPU. Ofrecía soporte de inferencia y entrenamiento para modelos de Hugging Face en TPUs. Ahora que TorchTPU trae soporte nativo, Optimum TPU ya no es necesario para ejecutar modelos de Hugging Face con Transformers y Diffusers.
Llega TorchTPU 🚀
TorchTPU añade TPU como un nuevo backend a PyTorch. Escribes código PyTorch normal. Viene con kernels de Pallas ajustados para TPUs, soporte para entrenamiento distribuido y una torch.compile integración para cuando quieres más velocidad. Instalarlo en una TPU VM requiere un solo comando:
pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu
El índice adicional te da la compilación de CPU de torch. Sin él, pip descarga la compilación de CUDA, que es mucho más grande y no necesitas en una TPU.
Tus bibliotecas favoritas, ahora en TPU
Transformers
Transformers v5 se construyó en torno a una idea: ser la fuente de verdad para las definiciones de modelos, de modo que puedan ejecutarse en cualquier lugar. «Cualquier lugar» ahora incluye las TPUs. Para cargar un modelo en una TPU, pasa device_map="tpu". Eso es todo.
from transformers import AutoModelForCausalLM, AutoTokenizer
model_id = "Qwen/Qwen3.5-9B"
model = AutoModelForCausalLM.from_pretrained(
model_id,
device_map="tpu", # will load model directly on the 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))
El modelo se ejecuta en modo eager por defecto. También puedes usar torch.compile para obtener mejor rendimiento, o ejecutarlo en una configuración distribuida.
¿Imágenes? ¿Videos? Diffusers los pinta en TPUs 🎨
TPU ya es un dispositivo de primera clase en Diffusers. Para prototipar o depurar, el modo eager solo necesita pipe.to("tpu"). Para despliegue, el modo compilado añade una única llamada compile . Tus pipelines y tu código permanecen igual, y no hay que mantener ningún fork de TPU.
import torch
from diffusers import StableDiffusion3Pipeline
pipe = StableDiffusion3Pipeline.from_pretrained(
"stabilityai/stable-diffusion-3.5-medium", dtype=torch.bfloat16
)
pipe.to("tpu") # eager mode: that's the whole change
# optional: compiled mode
pipe.transformer.compile(backend="tpu", fullgraph=True, dynamic=False)
image = pipe("a corgi astronaut planting a flag on a pizza moon").images[0]
La primera llamada compila, y cada llamada posterior con el mismo tamaño de imagen reutiliza el grafo.
¿Modelos demasiado grandes para un solo chip? Pasa un TensorParallelConfig a from_pretrained y el transformer se particiona entre los chips a medida que se carga, de modo que el modelo completo nunca tiene que caber en un solo lugar. Muchos pipelines populares lo admiten de fábrica, y los que no lo hacen están a solo una tp_plan de distancia. La guía de TPU de Diffusers tiene la receta completa.
Aquí hay algunos ejemplos ejecutados en una Cloud TPU v6e-4:
TP=4 (cada modelo dividido entre cuatro chips) · bf16 · batch 1
| MiniMax-H3 (video + audio) | Qwen-Image-2.1 | |
|---|---|---|
| "La mascota de peluche esponjosa de Hugging Face de la imagen, una cara amarilla y sonriente con mejillas rosadas y dos pequeñas manos abiertas, se yergue en la proa de un trirreme griego antiguo […] Sonido de olas rompiendo, madera crujiendo, viento en las jarcias, gaviotas lejanas y un swell orquestal épico y grave." · 1344×768, 5.2 s, 49 pasos · 188 s compilado, 262 s eager | "Un capibara con sombrero de mago, leyendo un libro a la luz de las velas, pintura al óleo" · 1024×1024, 40 pasos · 3.5 s compilado, 8.0 s eager | "Un letrero de neón de una tienda que dice 'QWEN IMAGE 2.1', noche lluviosa, reflejos sobre el pavimento mojado" · 2048×2048, 40 pasos |
Kernels
El Kernels Hub tiene kernels de TPU listos para usar. Carga uno con la librería Kernels y dale a tu modelo un impulso de velocidad.
Los kernels de TPU se escriben en Pallas, lo que te da control de bajo nivel sobre el hardware. También puedes incorporar kernels de la biblioteca Tokamax . Aquí tienes un kernel ragged-dot de Tokamax conectado a un modelo Mixture-of-Experts:
from kernels import get_kernel
from transformers import AutoModelForCausalLM
from transformers.integrations.moe import ALL_EXPERTS_FUNCTIONS
repo_id = "tengomucho/ragged-dot-tpu"
kernel = get_kernel(repo_id, version=1, trust_remote_code=[repo_id])
ALL_EXPERTS_FUNCTIONS.register("tokamax_ragged_dot", kernel.experts_forward)
model = AutoModelForCausalLM.from_pretrained(
"Qwen/Qwen3-30B-A3B-Instruct-2507",
experts_implementation="tokamax_ragged_dot",
)
Cuanto mayor sea el lote, mayor será la ganancia. Con un tamaño de lote de 64, la decodificación es 21.5× más rápida en Qwen3-30B-A3B y 12.8× más rápida en Gemma 4 26B-A4B.
Rendimiento de decodificación en una TPU v6e-8, kernel del Hub frente a la implementación predeterminada.
Lo probamos para que tú no tengas que hacerlo 🧪
Estamos mejorando continuamente la compatibilidad con TPU en transformers ejecutando la suite de pruebas de la biblioteca. Estamos trabajando para añadir un CI nocturno en TPU y hacer que los resultados estén disponibles en el (transformers dashboard)[https://huggingface.co/spaces/transformers-community/transformers-ci-dashboard].
El dashboard de CI de Transformers, con los resultados de TPU junto a los de otros aceleradores. ¡Trabajo en curso!
Todavía no está todo en verde, y eso nos parece bien. De las pruebas que se ejecutan en TPU, pasa el 96.5%. Cada fallo que encontramos se convierte en un pequeño script de reproducción y en un issue en el repositorio de TorchTPU. Hasta ahora hemos reportado 34 issues, y 19 ya están corregidos. Ejecutar en TPU también reveló algunos errores en el propio Transformers, como en StaticCache y en los planes tensor-parallel de Qwen3. Los corregimos upstream.
Si algo se rompe en TPU, queremos averiguarlo antes que tú.
Empieza ya
Las TPU ya no son una misión secundaria para los usuarios de PyTorch. Instala TorchTPU, elige un modelo del Hub y envíalo a "tpu".
Estamos deseando ver qué construyes 🤗. Si tienes algún problema, abre un issue y cuéntanoslo.