Hugging Face anuncia la compatibilidad nativa de las TPU, los chips de aceleración diseñados por Google, en su ecosistema. Gracias al backend TorchTPU para PyTorch, Transformers, Diffusers y Kernels funcionan directamente en estos procesadores, según la publicación de la empresa.
- Hugging Face permite usar las TPU de forma nativa mediante el backend TorchTPU para PyTorch.
- Transformers, Diffusers y Kernels ya funcionan directamente en TPU, y Optimum TPU deja de ser necesario para Transformers y Diffusers.
- TorchTPU incluye núcleos Pallas, entrenamiento distribuido e integración con torch.compile, y se instala con un solo comando.
- Las pruebas de Transformers en TPU alcanzan un 96,5 % de éxito, y Hugging Face ya ha corregido varios errores detectados.
Una TPU, o Unidad de Procesamiento de Tensores, es un chip diseñado para acelerar los cálculos matriciales que ocupan la mayor parte del tiempo de las redes neuronales. Google desarrolla estos chips desde hace una década y los utiliza para entrenar y explotar sus propios modelos, incluido Gemini. Están disponibles en Google Cloud, desde un solo chip hasta conjuntos de varios miles.
Hasta ahora, el ecosistema de software de las TPU se había construido alrededor de TensorFlow y JAX, y los usuarios de PyTorch debían realizar un mayor esfuerzo. Para ejecutar PyTorch en TPU, era necesario utilizar PyTorch/XLA (torch_xla), una solución que la publicación describe como funcional pero poco práctica: los tensores se evaluaban de forma perezosa, por lo que el código no siempre se ejecutaba en el momento esperado, y pequeñas modificaciones podían desencadenar recompilaciones del grafo. Hugging Face había desarrollado Optimum TPU para ofrecer inferencia y entrenamiento de sus modelos en TPU. Con TorchTPU, esta herramienta ya no es necesaria para ejecutar los modelos de Transformers y Diffusers.
TorchTPU, un backend nativo para PyTorch
TorchTPU añade la TPU como un nuevo backend de PyTorch, sin modificar el código habitual de PyTorch. El backend incluye núcleos Pallas optimizados para TPU, compatibilidad con entrenamiento distribuido e integración con torch.compile, para los usuarios que buscan mayor velocidad. Su instalación en una máquina virtual TPU se realiza con un solo comando: pip install torch-tpu –extra-index-url https://download.pytorch.org/whl/cpu. La publicación precisa que este índice proporciona la versión CPU de torch, más ligera que la versión CUDA que pip descarga por defecto y que una TPU no necesita.
Transformers, Diffusers y Kernels en TPU
En Transformers v5, cuyo objetivo es ser la referencia para las definiciones de modelos para que funcionen en todas partes, la carga de un modelo en TPU se realiza mediante el parámetro device_map= »tpu ». La publicación ilustra este paso con el modelo Qwen/Qwen3.5-9B. Por defecto, el modelo se ejecuta en modo eager, pero también puede utilizar torch.compile o una configuración distribuida.
En Diffusers, la TPU es ahora un dispositivo de primer nivel. El modo eager, adecuado para prototipado, solo requiere una llamada a pipe.to(« tpu »), mientras que el modo compilado, destinado al despliegue, añade una llamada de compilación. Para los modelos demasiado grandes para un solo chip, un TensorParallelConfig pasado a from_pretrained distribuye el transformador entre varios chips durante la carga. La publicación presenta ejemplos ejecutados en un Cloud TPU v6e-4, con los modelos MiniMax-H3, para vídeo y audio, y Qwen-Image-2.1.
La biblioteca Kernels da acceso a núcleos TPU listos para usar, escritos en Pallas, que ofrece un control de bajo nivel sobre el hardware. También se pueden integrar núcleos de la biblioteca Tokamax. Según la publicación, para un lote de 64, la decodificación es 21,5 veces más rápida en Qwen3-30B-A3B y 12,8 veces más rápida en Gemma 4 26B-A4B, en comparación con la implementación por defecto, con una TPU v6e-8.
Pruebas continuas
Hugging Face indica que ejecuta continuamente la suite de pruebas de Transformers en TPU y trabaja en añadir una integración continua nocturna en estos chips, cuyos resultados se publicarían en un panel dedicado. Según la publicación, el 96,5 % de las pruebas ejecutadas en TPU tienen éxito. Cada fallo da lugar a un script de reproducción y a un informe en el repositorio TorchTPU: se han reportado 34 problemas, de los cuales 19 ya están corregidos. Estas pruebas también han revelado errores en Transformers, concretamente en StaticCache y en los planes de paralelismo tensorial de Qwen3, que se han corregido de forma anticipada. En caso de problema, se invita a los usuarios a abrir un issue.
