Hugging Face kondigt native ondersteuning aan voor TPU’s, de versnellingschips ontworpen door Google, in zijn ecosysteem. Dankzij de backend TorchTPU voor PyTorch werken Transformers, Diffusers en Kernels rechtstreeks op deze processors, volgens het bericht van het bedrijf.
Een TPU, of Tensor Processing Unit, is een chip ontworpen om matrixberekeningen te versnellen die het grootste deel van de tijd van neurale netwerken in beslag nemen. Google ontwikkelt deze chips sinds ongeveer tien jaar en gebruikt ze om zijn eigen modellen te trainen en uit te voeren, waaronder Gemini. Ze zijn beschikbaar op Google Cloud, van een enkele chip tot clusters van enkele duizenden.
Tot nu toe was de software-ecosysteem van TPU’s gebouwd rond TensorFlow en JAX, en moesten gebruikers van PyTorch meer moeite doen. Om PyTorch op TPU uit te voeren, moest men gebruik maken van PyTorch/XLA (torch_xla), een oplossing die het bericht beschrijft als functioneel maar onhandig: tensoren werden lui geëvalueerd, zodat de code niet altijd op het verwachte moment werd uitgevoerd, en kleine wijzigingen konden hercompilaties van de grafiek veroorzaken. Hugging Face had Optimum TPU ontwikkeld om inferentie en training van zijn modellen op TPU mogelijk te maken. Met TorchTPU is deze tool niet langer nodig om de modellen van Transformers en Diffusers te laten draaien.
TorchTPU, een native backend voor PyTorch
TorchTPU voegt de TPU toe als een nieuwe backend voor PyTorch, zonder de gebruikelijke PyTorch-code te wijzigen. De backend bevat Pallas-kernels geoptimaliseerd voor TPU, ondersteuning voor gedistribueerde training en integratie met torch.compile, voor gebruikers die meer snelheid zoeken. De installatie op een TPU-virtuele machine gebeurt met één enkele opdracht: pip install torch-tpu –extra-index-url https://download.pytorch.org/whl/cpu. Het bericht preciseert dat deze index de CPU-versie van torch levert, lichter dan de CUDA-versie die pip standaard downloadt en waar een TPU geen behoefte aan heeft.
Transformers, Diffusers en Kernels op TPU
In Transformers v5, waarvan het doel is de referentie te zijn voor modeldefinities zodat ze overal werken, verloopt het laden van een model op TPU via de parameter device_map= »tpu ». Het bericht illustreert deze stap met het model Qwen/Qwen3.5-9B. Standaard wordt het model uitgevoerd in eager-modus, maar het kan ook torch.compile of een gedistribueerde configuratie gebruiken.
In Diffusers is de TPU nu een eersteklas apparaat. De eager-modus, geschikt voor prototyping, vereist slechts een aanroep naar pipe.to(« tpu »), terwijl de gecompileerde modus, bedoeld voor implementatie, een compilatie-aanroep toevoegt. Voor modellen die te groot zijn voor één enkele chip, verdeelt een TensorParallelConfig doorgegeven aan from_pretrained de transformer over meerdere chips bij het laden. Het bericht presenteert voorbeelden uitgevoerd op een Cloud TPU v6e-4, met de modellen MiniMax-H3, voor video en audio, en Qwen-Image-2.1.
De bibliotheek Kernels geeft toegang tot kant-en-klare TPU-kernels, geschreven in Pallas, dat een laag niveau van controle over de hardware biedt. Kernels afkomstig van de Tokamax-bibliotheek kunnen ook worden geïntegreerd. Volgens het bericht is het decoderen voor een batch van 64, 21,5 keer sneller op Qwen3-30B-A3B en 12,8 keer sneller op Gemma 4 26B-A4B, vergeleken met de standaardimplementatie, met een TPU v6e-8.
Doorlopende tests
Hugging Face geeft aan continu de testsuite van Transformers op TPU uit te voeren en te werken aan het toevoegen van een nachtelijke continue integratie op deze chips, waarvan de resultaten zouden worden gepubliceerd op een speciaal dashboard. Volgens het bericht slagen 96,5% van de uitgevoerde tests op TPU. Elk falen leidt tot een reproductiescript en een melding in de TorchTPU-opslagplaats: 34 problemen zijn gerapporteerd, waarvan 19 al zijn opgelost. Deze tests hebben ook bugs in Transformers aan het licht gebracht, met name in StaticCache en in de tensorparallelisme-plannen van Qwen3, die upstream zijn gecorrigeerd. Bij problemen worden gebruikers uitgenodigd een issue te openen.
