Google publica TorchTPU como código abierto y lleva PyTorch nativo a sus TPU

Google anunció el 20 de octubre de 2026, en PyTorch Conference 2026, que TorchTPU, su pila nativa de PyTorch para unidades de procesamiento tensorial (TPU), ya está disponible como código abierto en GitHub, según informó la empresa en el blog de Google Cloud. El software se desarrolló conjuntamente con el equipo de PyTorch de Meta y está diseñado para ejecutar cargas de trabajo de PyTorch existentes en TPU con cambios mínimos en el código.

Con este lanzamiento, las TPU se suman a las GPU y las CPU como opción nativa para PyTorch, según informó Google, y abarcan entrenamiento, inferencia y despliegue mediante herramientas como vLLM, SGLang, llm-d, TorchTitan, Ray, TorchRec, Hugging Face y LiteRT. Los desarrolladores pueden depurar modelos en modo eager y, después, alcanzar un rendimiento de nivel de producción mediante torch.compile para optimizar el uso de las TPU y reducir la sobrecarga de memoria.

Jana van Greunen, directora de Ingeniería de PyTorch en Meta, afirmó en el anuncio que TorchTPU permite a Meta desarrollar hardware de forma flexible junto con su equipo central de sistemas de aprendizaje automático y, al mismo tiempo, dar soporte a cargas de trabajo importantes de LLM, anuncios y sistemas de recomendación. «Nos permite desplegar nuestras cargas de trabajo de PyTorch en las TPU de Google, ofrecer comentarios directamente a los equipos de ingeniería y crear la flexibilidad del ecosistema que todos necesitan», señaló.

Integraciones con Hugging y de inferencia

Google se asoció con Hugging Face para que los modelos de pesos abiertos disponibles en Hugging Face Hub se ejecuten de forma nativa en TorchTPU. En Transformers v5, desplegar un modelo en TPU se hace pasando device_map="tpu", y la biblioteca Diffusers se ejecuta de forma nativa en TPU y admite el modo eager, torch.compile y el paralelismo de tensores entre varios aceleradores. Google afirmó que también creó canalizaciones de CI/CD específicas para probar de forma continua, cada noche, la suite Transformers de Hugging Face en la infraestructura de TPU. Jeff Boudier, vicepresidente de Producto de Hugging Face, dijo que esta colaboración lleva a la práctica el compromiso de la plataforma con la posibilidad de elegir el hardware.

Para la inferencia, el backend nativo de PyTorch de vLLM en TorchTPU se desarrolló en colaboración con Inferact. Google afirmó que admite el procesamiento por lotes continuo, la cuantización y el paralelismo entre varios chips en hardware TPU mediante kernels Pallas personalizados, a la vez que conserva los modelos de trabajadores, los planificadores y los servidores de API del proyecto original, por lo que las pilas de inferencia existentes no requieren cambios. Simon Mo, director ejecutivo y cofundador de Inferact, afirmó que la integración incorpora compatibilidad de primer nivel con vLLM en las TPU.

Un backend nativo de PyTorch para SGLang, desarrollado junto con RadixArk, está dirigido a cargas de trabajo estructuradas y basadas en agentes. Según el anuncio, SGLang utiliza RadixAttention para reutilizar la caché KV entre solicitudes, de modo que se omite el recálculo de los prefijos de instrucciones repetidos; además, la versión de TorchTPU emplea un backend de transferencia desagregado que también sirve como backend de E/S para HiCache. Banghua Zhu, director de tecnología y cofundador de RadixArk, afirmó que TorchTPU ha acelerado el trabajo de la empresa a medida que incorpora compatibilidad con SGLang en distintos tipos de silicio.

Google también lanza llm-d 1.0 junto con Red Hat, IBM y una comunidad de más de 500 colaboradores y 25 organizaciones, como proyecto Sandbox de la Cloud Native Computing Foundation. El anuncio describe llm-d como un vínculo entre el ecosistema de inferencia de PyTorch y el despliegue de producción basado en Kubernetes: su enrutador utiliza enrutamiento consciente de la caché de prefijos para enviar cada solicitud a la réplica que ya tiene su contexto y separa la fase de prefill, que exige mucho cómputo, de la fase de decode, que consume mucha memoria. La versión 1.0 incorpora control del flujo multiinquilino, alta disponibilidad del enrutador, API v1 estables y enrutamiento con reconocimiento de sesiones para agentes; además, en Google Kubernetes Engine, un endpoint común de llm-d puede ejecutarse tanto en GPU como en TPU.

El profesor Sashikumaar Ganesan, fundador de Zenteiq, afirmó en el anuncio que GKE Inference Gateway, impulsado por llm-d, ofreció un rendimiento agregado de inferencia más de 4× superior para la generación de datos sintéticos de la empresa, y elevó la capacidad diaria de generación de 1.470 millones a 5.970 millones de tokens por cada 1.000 aceleradores TPU v6e.

Entrenamiento y orquestación

TorchTitan, la implementación de referencia nativa de PyTorch para el entrenamiento distribuido, ahora incluye una integración de TorchTPU de nivel de producción, según la describe Google, basada en torch.compile, SimpleFSDP de TorchTitan, kernels Pallas especializados y optimizaciones a nivel de grafo. TorchTitan ofrece recetas de TorchTPU para preentrenar modelos densos, como Llama 3 70B y 405B, Qwen 3 8B y 32B, y MuseGlimmer 30B, además del modelo de mezcla de expertos Qwen 3 235B y el transformador de difusión FLUX.1 12B. Google afirmó que la integración se ha validado con hasta 8.000 TPU Ironwood (TPU v7x) y que el entrenamiento de TorchTPU mediante TorchTitan alcanzó una utilización de FLOPS del modelo de hasta el 52 % con Qwen 3 8B, con un escalado casi lineal en hasta 8.192 TPU Ironwood.

Google también incorporó a VeRL, un marco de aprendizaje por refuerzo, un backend para TPU y una receta de referencia integral. La receta GRPO, basada en una arquitectura desagregada de entrenador y muestreador y demostrada con conjuntos de datos de razonamiento matemático, está diseñada para que los desarrolladores ejecuten cargas de trabajo de VeRL en TPU y conserven su código existente de modelos y algoritmos, independiente del hardware.

Ray incorpora compatibilidad nativa con TorchTPU para coordinar el procesamiento de datos, el entrenamiento, el aprendizaje por refuerzo y la puesta en servicio de PyTorch en CPU, GPU y TPU dentro de un único flujo de trabajo. Edward Oakes, arquitecto jefe de Anyscale, afirmó que las empresas trabajaron para que las TPU fueran aceleradores de primera clase en Ray y añadió: «TorchTPU integra las TPU de Google directamente en esa pila, estándar del sector».

Arquitectura anterior y disponibilidad

TorchTPU se presentó públicamente el 7 de abril de 2026, en una publicación del blog de Google Developers que describía un diseño «Eager First» basado en la interfaz PrivateUse1 de PyTorch, con modos de ejecución anticipada de depuración, estricta y fusionada, respaldados por una caché de compilación compartida. En aquel momento, Google afirmó que Fused Eager ofrece una mejora del rendimiento del 50 % a más del 100 % con respecto a Strict Eager, y que torch.compile en TorchTPU captura gráficos FX con Torch Dynamo y los transforma mediante XLA y StableHLO, con kernels personalizados escritos en Pallas y JAX. La publicación señalaba que la pila admite Distributed Data Parallel, FSDPv2 y DTensor de forma inmediata, y que su arquitectura está diseñada para gestionar ejecuciones divergentes entre ranks. También incluía el lanzamiento de un repositorio público de GitHub entre los elementos de la hoja de ruta de TorchTPU para 2026.

El repositorio torch_tpu ya es público en la organización google-pytorch de GitHub y está publicado bajo la licencia Apache 2.0. Su archivo README enumera kernels ATen personalizados, una caché de compilación para cálculos de XLA, compatibilidad con entrenamiento distribuido mediante paralelismo de datos y de tensores, integración con torch.compile, utilidades de depuración y evaluación comparativa, y ejemplos para Llama, Qwen, ResNet, DLRM, minGPT y UNet. El proyecto es compatible con Python 3.11 y versiones posteriores, y el paquete torch-tpu 0.2.0 se publicó en PyPI el 19 de octubre de 2026, con paquetes wheel para Linux x86-64 y CPython de las versiones 3.11 a 3.14. Hay una guía del usuario en torchtpu.readthedocs.io, y el archivo README indica que los cuadernos interactivos del proyecto deben ejecutarse desde una máquina virtual TPU; una instancia v6e de GCE es la opción más sencilla.

En el mismo anuncio, Google afirmó que LiteRT, antes conocido como TensorFlow Lite, se incorpora al ecosistema de PyTorch: la biblioteca LiteRT-Torch convierte un modelo de PyTorch con una sola litert-torch.convert() llamada, y los modelos convertidos se ejecutan en Android, iOS, macOS, Windows, Linux, la web y dispositivos IoT. Un conjunto de agentes para TPU, pensado para automatizar la incorporación, el diagnóstico y el ajuste del rendimiento de inferencia de modelos de PyTorch, está disponible en versión preliminar inicial en GitHub.

El anuncio remite a los desarrolladores a la documentación de TPU de Google Cloud para empezar a ejecutar PyTorch en las TPU.