mrkeyoor.com_
Tue 20 Oct 21:05 UTC
Open Source3 min read

TorchTPU Brings Native TPU Support to Hugging Face Libraries

Transformers and Diffusers now use TPUs through native PyTorch. Hugging Face reports a 96.5% pass rate for the tests it runs on TPU.

Hugging Face announced native TPU support for Transformers and Diffusers on October 20. The integration uses TorchTPU, an open-source PyTorch backend, so developers can try Google's accelerators while keeping their model code in PyTorch. Compatibility remains a work in progress: Hugging Face reports a 96.5% pass rate for the tests it runs on TPU.

TorchTPU was developed by Google and the PyTorch team at Meta, according to its version 0.2.0 documentation. It integrates with PyTorch without JAX runtime dependencies. This article examines the announcement and those versioned documents; we haven't installed the backend or run a TPU workload.

What changes in the model code

Hugging Face's Transformers example loads Qwen3.5-9B with device_map="tpu". Its Diffusers example moves a pipeline with pipe.to("tpu") and offers an optional compilation call. Hugging Face says its separate Optimum TPU package is no longer required to run these libraries on TPUs. These are publisher-provided examples, rather than tests performed by MrKeyoor.

The backend's migration instructions describe changing device="cuda" to device="tpu". Distributed workloads also change the communication backend from nccl to tpu_dist. That second change matters for anyone moving beyond one device: model placement and communication between devices are separate parts of a port.

The repository is licensed under Apache 2.0 and currently identifies 0.2.0 as its public package release. Its installation instructions request the CPU build of PyTorch through an extra package index, avoiding an unnecessary CUDA build. The package version identifies the documented setup here; it doesn't establish when each capability first appeared.

Eager execution still involves compilation

The execution guide explains what happens beneath the familiar interface. TorchTPU initially records an operation on the CPU, queues it on a stream, then prepares and compiles the work or retrieves a cached executable. In default eager mode, it processes operations individually and asynchronously. Runtime errors can therefore surface after the operation that caused them.

With torch.compile, the backend receives a larger computation graph and applies more optimization before execution, according to the same guide. For a migration trial, that makes cold-start time and performance after compilation separate measurements worth recording. The guide also requires synchronization when waiting for TPU work to finish, so timing only the Python dispatch would miss part of the work.

Variable sequence lengths deserve particular attention. The 0.2.0 introduction warns that language-model decoding can trigger recompilation at every step unless bounded dynamism is enabled. That mechanism lets a compiled graph accommodate changing sizes within specified bounds. A successful model load alone won't reveal whether generation repeatedly pays compilation costs.

What the pass rate covers

Hugging Face's test report limits its 96.5% figure to tests that run on TPU. It says the team reported 34 issues and that 19 were already fixed. The percentage doesn't establish that the entire Transformers suite, every Hub model or a reader's production workload passes. The announcement supplies no total test count alongside that percentage.

The next check is whether Hugging Face's planned nightly TPU testing produces regularly available results. Its announcement says the team is working to add that CI coverage and expose results in the Transformers dashboard. For developers considering a port, results for the model and operations they use will be more useful than an aggregate pass rate alone.

Sources

  1. Hugging Face: TorchTPU, TPUs Natively
  2. TorchTPU 0.2.0: Introduction
  3. TorchTPU 0.2.0: Execution model
  4. google-pytorch/torch_tpu repository