Google disponibiliza o TorchTPU como código aberto, levando o PyTorch nativo às suas TPUs

Em 20 de outubro de 2026, durante a PyTorch Conference 2026, o Google anunciou que o TorchTPU, sua pilha nativa de PyTorch para unidades de processamento tensorial (TPUs), já está disponível como código aberto no GitHub, informou a empresa no blog do Google Cloud. O software foi desenvolvido em conjunto com a equipe do PyTorch da Meta e foi projetado para executar cargas de trabalho existentes do PyTorch em TPUs com o mínimo de alterações no código.

Com o lançamento, as TPUs passam a ser uma opção nativa para PyTorch, ao lado de GPUs e CPUs, abrangendo treinamento, disponibilização e implantação por meio de ferramentas como vLLM, SGLang, llm-d, TorchTitan, Ray, TorchRec, Hugging Face e LiteRT, informou o Google. Os desenvolvedores podem depurar modelos no modo eager e, em seguida, alcançar desempenho de nível de produção com o torch.compile, para otimizar o uso das TPUs e reduzir a sobrecarga de memória.

Jana van Greunen, diretora de engenharia do PyTorch na Meta, afirmou no anúncio que o TorchTPU permite à Meta desenvolver hardware em conjunto com flexibilidade com sua equipe central de sistemas de aprendizado de máquina, dando suporte também a cargas de trabalho importantes de LLMs, anúncios e sistemas de recomendação. “Isso nos permite implantar nossas cargas de trabalho do PyTorch nas TPUs do Google, dar feedback diretamente às equipes de engenharia e criar a flexibilidade de ecossistema de que todos precisam”, afirmou.

Integrações com Hugging e para disponibilização de modelos

O Google fez uma parceria com Hugging Face para que os modelos de pesos abertos no Hugging Face Hub sejam executados nativamente no TorchTPU. No Transformers v5, a implantação de um modelo em TPUs é feita passando device_map="tpu", e a biblioteca Diffusers é executada nativamente em TPUs, com suporte ao modo eager, ao torch.compile e ao paralelismo de tensores em vários aceleradores. O Google afirmou também ter criado pipelines dedicados de CI/CD para testes noturnos contínuos da suíte Hugging Face Transformers na infraestrutura de TPU. Jeff Boudier, vice-presidente de produto da Hugging Face, disse que a colaboração põe em prática o compromisso da plataforma com a liberdade de escolha de hardware.

Para disponibilização de modelos, o backend nativo de PyTorch do vLLM no TorchTPU foi desenvolvido em parceria com a Inferact. Segundo o Google, ele oferece suporte a processamento contínuo em lotes, quantização e paralelismo entre vários chips no hardware TPU, usando kernels Pallas personalizados e preservando os modelos de workers, escalonadores e servidores de API do projeto upstream, de modo que as pilhas de inferência existentes não exijam alterações. Simon Mo, CEO e cofundador da Inferact, afirmou que a integração oferece suporte de primeira classe ao vLLM em TPUs.

Um backend nativo de PyTorch para SGLang, desenvolvido com a RadixArk, é voltado a cargas de trabalho estruturadas e agênticas. Segundo o anúncio, o SGLang usa o RadixAttention para reutilizar o cache KV entre solicitações, permitindo que prefixos de prompt repetidos evitem o recálculo; já a versão do TorchTPU emprega um backend de transferência desagregada que também funciona como backend de E/S para o HiCache. Banghua Zhu, CTO e cofundador da RadixArk, afirmou que o TorchTPU acelerou o trabalho da empresa enquanto ela amplia o suporte ao SGLang em diferentes plataformas de silício.

O Google também está lançando o llm-d 1.0 com a Red Hat, a IBM e uma comunidade de mais de 500 colaboradores e 25 organizações, como projeto Sandbox da Cloud Native Computing Foundation. O anúncio descreve o llm-d como uma ponte entre o ecossistema de serving do PyTorch e a implantação de produção baseada em Kubernetes: seu roteador usa roteamento ciente do cache de prefixos para enviar cada solicitação à réplica que já contém seu contexto, e separa o prefill, que exige muito processamento, do decode, que exige muita memória. A versão 1.0 adiciona controle de fluxo multi-inquilino, alta disponibilidade do roteador, APIs v1 estáveis e roteamento de agentes com reconhecimento de sessão; no Google Kubernetes Engine, um endpoint comum do llm-d pode ser executado tanto em GPUs quanto em TPUs.

O professor Sashikumaar Ganesan, fundador da Zenteiq, afirmou no anúncio que o GKE Inference Gateway, com tecnologia do llm-d, proporcionou uma vazão agregada de inferência mais de 4× maior para a geração de dados sintéticos da empresa, elevando a capacidade diária de geração de 1,47 bilhão para 5,97 bilhões de tokens por 1.000 aceleradores TPU v6e.

Treinamento e orquestração

O TorchTitan, implementação de referência nativa do PyTorch para treinamento distribuído, agora inclui uma integração do TorchTPU de nível de produção, descrita pelo Google, baseada em torch.compile, no SimpleFSDP do TorchTitan, em kernels Pallas especializados e em otimizações no nível do grafo. O TorchTitan oferece receitas do TorchTPU para pré-treinar modelos densos, como Llama 3 70B e 405B, Qwen 3 8B e 32B e MuseGlimmer 30B, além do modelo de mistura de especialistas Qwen 3 235B e do transformer de difusão FLUX.1 12B. O Google afirmou que a integração foi validada em até 8.000 TPUs Ironwood (TPU v7x) e que o treinamento do TorchTPU pelo TorchTitan alcançou até 52% de utilização dos FLOPs do modelo no Qwen 3 8B, com escalabilidade quase linear em até 8.192 TPUs Ironwood.

O Google também incorporou ao VeRL, um framework de aprendizado por reforço, um backend para TPU e uma receita de referência de ponta a ponta. A receita GRPO, baseada em uma arquitetura desagregada de treinador e amostrador e demonstrada em conjuntos de dados de raciocínio matemático, tem o objetivo de permitir que os desenvolvedores executem cargas de trabalho do VeRL em TPUs, mantendo o código existente de modelos e algoritmos independente de hardware.

Ray oferece suporte nativo ao TorchTPU para orquestrar o processamento de dados, o treinamento, o aprendizado por reforço e o serviço de inferência com PyTorch em CPUs, GPUs e TPUs, tudo em um único fluxo de trabalho. Edward Oakes, arquiteto-chefe da Anyscale, afirmou que as empresas trabalharam para tornar as TPUs aceleradores de primeira classe no Ray, acrescentando: “O TorchTPU integra diretamente as TPUs do Google a essa pilha padrão do setor.”

Arquitetura e disponibilidade anteriores

O TorchTPU foi apresentado publicamente em 7 de abril de 2026, em uma publicação no blog Google Developers que descrevia um projeto “Eager First”, baseado na interface PrivateUse1 do PyTorch, com modos de execução eager Debug, Strict e Fused apoiados por um cache de compilação compartilhado. Na ocasião, o Google afirmou que o Fused Eager oferece um aumento de desempenho de 50% a mais de 100% em relação ao Strict Eager e que o torch.compile no TorchTPU captura grafos FX com o Torch Dynamo e os converte por meio do XLA e do StableHLO, usando kernels personalizados escritos em Pallas e JAX. A publicação dizia que a pilha oferece suporte nativo a Distributed Data Parallel, FSDPv2 e DTensor e foi projetada para lidar com execuções divergentes entre ranks. Também listava o lançamento de um repositório público no GitHub entre os itens do roteiro do TorchTPU para 2026.

O repositório torch_tpu agora é público na organização google-pytorch do GitHub e está licenciado sob a Apache 2.0. O README lista kernels ATen personalizados, um cache de compilação para cálculos do XLA, suporte a treinamento distribuído com paralelismo de dados e de tensores, integração com torch.compile, utilitários de depuração e benchmarking, além de exemplos para Llama, Qwen, ResNet, DLRM, minGPT e UNet. O projeto é compatível com Python 3.11 e versões posteriores, e o pacote torch-tpu 0.2.0 foi publicado no PyPI em 19 de outubro de 2026, com pacotes wheel para Linux x86-64 e CPython 3.11 a 3.14. Um guia do usuário está disponível em torchtpu.readthedocs.io, e o README observa que os notebooks interativos do projeto precisam ser executados em uma VM TPU, sendo uma instância GCE v6e a opção mais simples.

No mesmo anúncio, o Google afirmou que o LiteRT, anteriormente chamado TensorFlow Lite, está se integrando ao ecossistema do PyTorch: a biblioteca LiteRT-Torch converte um modelo do PyTorch com uma única litert-torch.convert() chamada, e os modelos convertidos são executados em dispositivos Android, iOS, macOS, Windows, Linux, na web e em IoT. Uma suíte de agentes para TPU, destinada a automatizar a integração, o diagnóstico e o ajuste do desempenho de inferência de modelos do PyTorch, está em fase inicial de prévia no GitHub.

O anúncio direciona os desenvolvedores à documentação das TPUs do Google Cloud para começar a executar o PyTorch em TPUs.