Eigen RadarAI
Archive
Analysis

TorchTPU brings Google’s AI chips into familiar PyTorch workflows

Google has released TorchTPU, opening a native route for PyTorch applications to run on its AI accelerators. Hugging Face has added support in its model libraries, including Transformers and Diffusers. Developers can retain familiar code while choosing different hardware, with compilation and distributed execution among the supported paths.

Artificial Intelligence··Night
In an imagined machine workshop, differently sized model cassettes dock into a shared accelerator carrier extending behind them.

Google chips gain a native PyTorch route

Google has opened TorchTPU, software that lets applications built with PyTorch run on its Tensor Processing Units (TPUs), the company’s AI accelerators. Hugging Face, which develops model libraries and hosts models, has added native support in Transformers, Diffusers and Kernels. Developers can select a TPU as their device while retaining ordinary PyTorch code. The change connects a widely used model-development framework with Google’s computing hardware.[1], [3]

Until this release, PyTorch users reached these chips through PyTorch/XLA. Hugging Face describes that earlier route as deferring computations, with changes sometimes requiring the computation graph to be compiled again. Its own Optimum TPU package helped models run on that infrastructure. The company says that package is no longer required for the new native path through Transformers and Diffusers.[1]

Models can run directly or compile

Transformers loads a model onto the selected device and uses eager execution by default, carrying out operations without deferring them. TorchTPU also supports compilation and distributed training. Its TPU-tuned Pallas kernels provide the lower-level computation routines. Hugging Face’s Kernels library can load those routines, including implementations from Tokamax, for use with the model libraries.[1]

Diffusers offers eager execution for development and an optional compilation step for deployment. After compilation, repeated calls with the same image size reuse the computation graph. Models larger than a single chip’s capacity can be divided across chips during loading. Hugging Face demonstrates this route with image and video generation on Cloud TPU hardware.[1]

Migration tools accompany the release

Google says it worked with Meta, Hugging Face, Inferact and RadixArk to support established tools. Its TPU agent suite is intended to help move applications and fix migration errors. Hugging Face reports that some TPU tests still fail and that nightly testing and a public results dashboard remain in development. Hugging Face says it has reported 34 TorchTPU issues and that 19 have been fixed; nightly TPU tests remain under development.[1], [2]

References

  1. News sourceHugging FaceHugging Face brings native TPU support to its model libraries↩1↩2↩3↩4↩5
  2. News sourceTechCrunchGoogle releases TorchTPU for PyTorch applications on its AI chips↩
  3. News sourceGoogle CloudGoogle opens TorchTPU with native PyTorch support↩
AI safety and regulationChip wars and the supply chainThe frontier model race

More analysis from this desk

View all Brief analyses from the Artificial Intelligence desk →