O Google abriu o código do TorchTPU, uma camada de software que permite aos desenvolvedores executar PyTorch — um kit de ferramentas amplamente usado para criar e executar modelos de IA — nas unidades de processamento tensorial do Google, ou TPUs. Desenvolvido com a equipe do PyTorch da Meta, o projeto foi concebido para migrar cargas de trabalho existentes para o hardware TPU com mudanças relativamente pequenas no código, em vez de exigir uma reescrita.
Segundo o Google, o TorchTPU oferece suporte tanto à execução eager, que executa as operações à medida que são chamadas e é útil para depuração, quanto à execução compilada, voltada a cargas de trabalho de produção. Os desenvolvedores podem trocar a configuração do dispositivo de “cuda” para “tpu” e usar o backend distribuído de TPU; o Google também indica a API de detecção de aceleradores do PyTorch como alternativa. O TorchTPU está disponível no GitHub, mas o roteiro do projeto ainda prevê trabalho em formas de entrada dinâmicas, velocidade de compilação, cobertura de operadores, aprendizado por reforço e quantização.
O que muda para os desenvolvedores?
O lançamento aproxima as TPUs de ferramentas que muitos desenvolvedores de PyTorch já usam. Segundo o Google, o TorchTPU funciona com as ferramentas de inferência vLLM e SGLang, além de projetos de treinamento e orquestração, como TorchTitan, Ray e TorchRec. No Transformers v5, a Hugging Face diz que os desenvolvedores podem selecionar uma TPU com a configuração device_map="tpu". A biblioteca Diffusers também oferece suporte à execução em TPUs, inclusive no modo compilado e com a divisão do trabalho entre vários aceleradores.
Essa abordagem difere da adotada pelo PyTorch/XLA, caminho anterior do Google para executar PyTorch em TPUs. Segundo o Google, esse sistema exigia mudanças mais profundas no código, tinha tempos de compilação longos e não oferecia suporte ao modo eager do PyTorch. O Google criou um pipeline de testes que, segundo o InfoQ, avalia cerca de 5.000 modelos da Hugging Face em mais de 500 arquiteturas. Isso indica testes abrangentes, mas não significa que todos os modelos ou cargas de trabalho funcionarão sem alterações.
Como ele se compara a outras opções de hardware?
O Google afirma que o modo DeferAndFuse do TorchTPU — um híbrido de execução eager e compilada — pode ser pelo menos 50% mais rápido que o modo Strict Eager e reduzir o uso máximo de memória do acelerador. Em outro teste, o Google executou um modelo Muse Glimmer de 30 bilhões de parâmetros em PyTorch padrão: segundo o InfoQ, a taxa de inferência foi de 25.700 tokens por segundo em um chip TPUv7, com comprimento de sequência de 5.120, e de 22.400 com 10.240. O treinamento em quatro chips TPUv7 chegou a 5.600 e 4.600 tokens por segundo, respectivamente, nesses comprimentos. São resultados divulgados pelo Google, não uma comparação independente.
O TechCrunch compara o TorchTPU ao TorchNeuron, da Amazon, anunciado na mesma semana para executar cargas de trabalho do PyTorch nos chips Trainium da empresa. Ambos buscam facilitar o uso de alternativas às GPUs da Nvidia. O lançamento do TorchTPU acrescenta mais uma opção de hardware para quem usa PyTorch, mas não demonstra, por si só, que migrar uma carga de trabalho reduzirá custos ou terá desempenho superior ao da configuração atual de um desenvolvedor.
Por que importa
Para desenvolvedores que já usam PyTorch, o TorchTPU pode facilitar o teste ou a implantação de cargas de trabalho em TPUs do Google sem reconstruí-las do zero. Ainda não está claro se a migração melhora o custo ou a velocidade de um modelo específico; os números de desempenho citados são resultados do próprio Google.
