Google ha publicado como código abierto TorchTPU, una capa de software que permite a los desarrolladores ejecutar PyTorch —el conjunto de herramientas más usado para crear y ejecutar modelos de IA— en las unidades de procesamiento tensorial de Google, o TPU. Desarrollado junto con el equipo de PyTorch de Meta, está pensado para trasladar cargas de trabajo existentes al hardware TPU con cambios de código relativamente pequeños, en vez de tener que reescribirlas.
Google señala que TorchTPU admite tanto la ejecución eager, que ejecuta las operaciones a medida que se llaman y resulta útil para depurar, como la ejecución compilada para cargas de producción. Los desarrolladores pueden cambiar una opción del dispositivo de “cuda” a “tpu” y usar el backend distribuido para TPU; Google también señala que la API de PyTorch para detectar aceleradores es otra alternativa. TorchTPU está disponible en GitHub, aunque su hoja de ruta aún contempla mejoras en las formas dinámicas de entrada, la velocidad de compilación, la cobertura de operadores, el aprendizaje por refuerzo y la cuantización.
¿Qué cambia para los desarrolladores?
El lanzamiento conecta las TPU con software que muchos desarrolladores de PyTorch ya utilizan. Google afirma que TorchTPU funciona con las herramientas de servicio vLLM y SGLang, y con proyectos de entrenamiento y orquestación como TorchTitan, Ray y TorchRec. En Transformers v5, Hugging Face indica que los desarrolladores pueden seleccionar una TPU con la opción device_map="tpu". Su biblioteca Diffusers también admite la ejecución en TPU, incluido el modo compilado y la distribución del trabajo entre varios aceleradores.
Se trata de un enfoque distinto de PyTorch/XLA, la vía anterior de Google para ejecutar PyTorch en TPU. Google afirma que ese sistema exigía cambios de código más invasivos, tenía tiempos de compilación prolongados y no admitía el modo eager de PyTorch. Google ha creado una canalización de pruebas que, según InfoQ, evalúa unos 5.000 modelos de Hugging Face en más de 500 arquitecturas. Esto apunta a pruebas amplias, pero no significa que todos los modelos o cargas de trabajo vayan a funcionar sin cambios.
¿Cómo se compara con otras opciones de hardware?
Google afirma que el modo DeferAndFuse de TorchTPU —una combinación de ejecución eager y compilada— puede ser al menos un 50% más rápido que el modo Strict Eager y reducir, a la vez, el uso máximo de memoria del acelerador. En una prueba aparte, Google ejecutó un modelo Muse Glimmer de 30.000 millones de parámetros con PyTorch estándar: InfoQ informa de un rendimiento de inferencia de 25.700 tokens por segundo en un chip TPUv7 con una longitud de secuencia de 5.120, y de 22.400 con una longitud de 10.240. El entrenamiento en cuatro chips TPUv7 alcanzó 5.600 y 4.600 tokens por segundo, respectivamente. Son resultados comunicados por Google, no una comparación independiente.
TechCrunch sitúa TorchTPU junto a TorchNeuron, de Amazon, anunciado esa misma semana para ejecutar cargas de trabajo de PyTorch en los chips Trainium de Amazon. Ambos buscan facilitar el uso de alternativas a las GPU de Nvidia. El lanzamiento de TorchTPU suma otra opción de hardware para quienes usan PyTorch, pero por sí solo no demuestra que migrar una carga de trabajo vaya a reducir los costes o superar el rendimiento de la configuración actual de un desarrollador.
Por qué importa
Para los desarrolladores que ya usan PyTorch, TorchTPU podría facilitar las pruebas o la implementación de cargas de trabajo en las TPU de Google sin tener que reconstruirlas desde cero. No está claro si el cambio mejorará el coste o la velocidad de un modelo concreto; las cifras de rendimiento citadas son resultados de Google.
