Hugging Face anuncia o suporte nativo dos TPU, os chips de aceleração desenhados pela Google, no seu ecossistema. Graças ao backend TorchTPU para PyTorch, Transformers, Diffusers e Kernels funcionam diretamente nestes processadores, de acordo com o artigo publicado pela empresa.
Um TPU, ou Tensor Processing Unit, é um chip desenhado para acelerar os cálculos matriciais que ocupam a maior parte do tempo das redes neuronais. A Google desenvolve estes chips há cerca de dez anos e usa-os para treinar e explorar os seus próprios modelos, incluindo o Gemini. Estão disponíveis no Google Cloud, desde um único chip a conjuntos de vários milhares.
Até agora, o ecossistema de software dos TPUs tinha sido construído em torno do TensorFlow e do JAX, e os utilizadores do PyTorch tinham de fazer um esforço adicional. Para executar PyTorch num TPU, era necessário usar o PyTorch/XLA (torch_xla), uma solução que o artigo descreve como funcional mas pouco conveniente: os tensores eram avaliados de forma preguiçosa, pelo que o código nem sempre era executado no momento esperado, e pequenas modificações podiam desencadear recompilações do grafo. A Hugging Face tinha desenvolvido o Optimum TPU para oferecer inferência e treino dos seus modelos em TPU. Com o TorchTPU, esta ferramenta já não é necessária para executar os modelos de Transformers e Diffusers.
TorchTPU, um backend nativo para PyTorch
O TorchTPU adiciona o TPU como um novo backend do PyTorch, sem modificar o código PyTorch habitual. O backend inclui kernels Pallas otimizados para TPU, suporte para treino distribuído e integração com torch.compile, para os utilizadores que procuram mais velocidade. A sua instalação numa máquina virtual TPU faz-se com um único comando: pip install torch-tpu –extra-index-url https://download.pytorch.org/whl/cpu. O artigo especifica que este índice fornece a versão CPU do torch, mais leve do que a versão CUDA que o pip descarrega por padrão e de que um TPU não precisa.
Transformers, Diffusers e Kernels no TPU
No Transformers v5, cujo objetivo é ser a referência para as definições de modelos para que funcionem em qualquer lado, o carregamento de um modelo num TPU passa pelo parâmetro device_map= »tpu ». O artigo ilustra este passo com o modelo Qwen/Qwen3.5-9B. Por padrão, o modelo executa em modo eager, mas também pode usar torch.compile ou uma configuração distribuída.
No Diffusers, o TPU é agora um dispositivo de primeira linha. O modo eager, adequado para prototipagem, requer apenas uma chamada a pipe.to(« tpu »), enquanto o modo compilado, destinado à implementação, adiciona uma chamada de compilação. Para modelos demasiado grandes para um único chip, um TensorParallelConfig passado a from_pretrained distribui o transformador por vários chips durante o carregamento. O artigo apresenta exemplos executados num Cloud TPU v6e-4, com os modelos MiniMax-H3, para vídeo e áudio, e Qwen-Image-2.1.
A biblioteca Kernels dá acesso a kernels TPU prontos a usar, escritos em Pallas, que oferece um controlo de baixo nível sobre o hardware. Kernels provenientes da biblioteca Tokamax também podem ser integrados. Segundo o artigo, para um batch de 64, a descodificação é 21,5 vezes mais rápida no Qwen3-30B-A3B e 12,8 vezes mais rápida no Gemma 4 26B-A4B, em comparação com a implementação padrão, com um TPU v6e-8.
Testes contínuos
A Hugging Face indica que executa continuamente a suite de testes do Transformers no TPU e que está a trabalhar para adicionar uma integração contínua noturna nestes chips, cujos resultados seriam publicados num painel dedicado. Segundo o artigo, 96,5 % dos testes executados no TPU são bem-sucedidos. Cada falha dá origem a um script de reprodução e a um relatório no repositório TorchTPU: 34 problemas foram reportados, dos quais 19 já estão corrigidos. Estes testes também revelaram bugs no Transformers, nomeadamente no StaticCache e nos planos de paralelismo tensorial do Qwen3, que foram corrigidos a montante. Em caso de problema, os utilizadores são convidados a abrir uma issue.
