TorchTPU's DeferAndFuse can run at least 50% faster than Strict Eager

Top News · Reported · pytorch.org

TorchTPU is Google and Meta's open-source backend for running PyTorch natively on Google TPUs. Its DeferAndFuse mode can run at least 50% faster than Strict Eager and reduce peak HBM usage through fusion.

Here's what you can use:

• Changing model.to("cuda") to model.to("tpu") is the basic migration path.
• torch.compile gets a native TPU backend alongside eager execution, with support for approximately 450 Aten ops.
• The team validated out-of-the-box execution on the 5,000 most-downloaded Hugging Face models, spanning language, audio, vision, and text-encoder/search categories.
• Validation also covers DTensor sharding, multi-host inference with vLLM and SGLang, and TorchTitan training with Ray orchestration.

One catch: default Skinny Collectives can impose a 30-40% end-to-end performance penalty because XLA cannot overlap computation and communication.

Try it: On a TPU machine with the TorchTPU wheel installed, run compiled_model = torch.compile(model, backend="tpu").

Sources: pytorch.org