TorchTPU aproxima a TPU do código PyTorch existente
O backend usa o dispositivo tpu em Transformers e Diffusers; compatibilidade do fluxo completo continua sendo o ponto a conferir.
device_map="tpu" é o ponto de entrada mostrado pela Hugging Face para carregar um modelo do Transformers com TorchTPU. Em uma VM com TPU, o artigo documenta a instalação:
pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu
O índice adicional fornece a distribuição CPU do PyTorch, evitando baixar a distribuição CUDA nesse ambiente. O anúncio foi publicado em 20 de outubro de 2026.
TorchTPU oferece um backend nativo para PyTorch. O texto apresenta execução imediata, integração com torch.compile e suporte a treinamento distribuído. Para Diffusers, mostra a movimentação do pipeline com pipe.to("tpu"), seguida de compilação opcional.
O ganho arquitetural que eu priorizaria é manter a definição do modelo no fluxo habitual do projeto. Isso não dispensa conferir as operações utilizadas: os próprios autores relatam falhas ainda abertas na execução dos testes em TPU.
Minha escolha seria começar pelo modo imediato, validar a saída e só depois comparar a versão compilada. O limite de engenharia é a compatibilidade do seu pipeline completo, não apenas conseguir mover seus tensores para o dispositivo.
Fonte: Alvaro e coautores, da Hugging Face, em Hugging Face ❤️ TorchTPU: TPUs, Natively.