From the source
Hugging Face says TorchTPU, a native PyTorch backend, now lets Transformers, Diffusers and Kernels run on Google TPUs without porting code to torch_xla or learning JAX.
- Until now, running PyTorch on TPUs meant going through PyTorch/XLA, where tensors were lazy, small changes could trigger graph recompilations and debugging was hard to reason about.
- TorchTPU adds TPU as a new PyTorch backend with Pallas kernels tuned for TPUs, distributed training support and a torch.compile integration, installed on a TPU VM with one pip command.
- Loading a model on a TPU takes device_map="tpu" in Transformers, while Diffusers needs pipe.to("tpu") in eager mode or one compile call for deployment.
- Models too large for one chip can be sharded across chips by passing a TensorParallelConfig to from_pretrained.
- Hugging Face says Optimum TPU is no longer required for Hugging Face models with Transformers and Diffusers, and that 96.5% of the tests that run on TPU pass, with 34 issues reported to the TorchTPU repo and 19 already fixed.
