Google y Meta lanzan TorchTPU para ejecutar PyTorch de forma nativa en TPU
El backend de código abierto, iniciado en Google a mediados de 2025 junto con el equipo de PyTorch de Meta, ya se usa en producción por algunos de los mayores clientes de Google Cloud, según su equipo.
Ilustración generada con IAGoogle y Meta han presentado TorchTPU, un backend de código abierto que permite ejecutar PyTorch de forma nativa en los chips TPU de Google. El equipo del proyecto afirmó que comenzó en Google a mediados de 2025 en colaboración con el equipo de PyTorch de Meta, y que se encuentra en desarrollo activo.
Según el equipo, TorchTPU ya lo utilizan en producción, para cargas de trabajo de inferencia y entrenamiento, algunos de los mayores clientes de Google Cloud. El equipo señaló que, históricamente, las TPU han estado ligadas a los ecosistemas de TensorFlow y JAX.
El equipo describió TorchTPU como una ruptura con los intentos anteriores. Google lanzó en 2020 PyTorch/XLA, una primera versión compilada de PyTorch para TPU basada en Lazy Tensors. El equipo explicó que esa integración requería cambios intrusivos en el código fuente, sufría una latencia de compilación severa, carecía de compatibilidad con el modo de ejecución eager de PyTorch y iba rezagada en las funciones más recientes de computación distribuida. TorchTPU, en cambio, trata la ejecución eager como una prioridad máxima y no tiene dependencias de JAX en su entorno de ejecución.
El objetivo declarado es que trasladar una carga de trabajo a una TPU requiera poco más que cambiar en el código el dispositivo de destino de una GPU a una TPU. El equipo advirtió de que ese cambio por sí solo no ofrecería el máximo rendimiento de la TPU, porque las formas, el sharding y las operaciones de un modelo pueden estar optimizados para GPU, pero afirmó que el modelo debería funcionar igualmente listo para usar.
Hasta ahora, según el equipo, el proyecto ha:
- añadido compatibilidad con los modos eager y compilado, incluido un nuevo modo eager llamado DeferAndFuse;
- añadido compatibilidad con unas 450 operaciones de ATen;
- habilitado kernels personalizados escritos mediante JAX/Pallas y el DSL de Helion;
- validado la ejecución lista para usar en los 5,000 modelos más descargados de Hugging Face;
- validado la integración con vLLM, SGLang, TorchTitan, Ray y TorchRec.
El equipo también enumeró limitaciones. Un dispositivo TPU solo puede ser accedido por un proceso de espacio de usuario, por lo que la colocalización de aprendizaje por refuerzo no puede admitirse de forma predeterminada, y el equipo afirmó que investiga formas de sortearlo. También está añadiendo compatibilidad con dinamismo acotado, ya que el compilador y el entorno de ejecución subyacentes asumen que las formas de los tensores se conocen en tiempo de compilación.
En cuanto a la precisión, TorchTPU usa por defecto un modo medio que multiplica en bfloat16 y acumula en FP32. Los modos superiores recuperan más precisión FP32 con un costo de cómputo de aproximadamente 3 o 6 veces, y el equipo recomienda el modo predeterminado para la mayoría de las cargas de trabajo de aprendizaje profundo.