Google开源TorchTPU，原生支持PyTorch运行于Cloud TPU
Google Open-Sources TorchTPU, a Native PyTorch Backend for Cloud TPUs
TPU用户必看，TorchTPU正式取代PyTorch/XLA成为官方推荐后端，完整覆盖训练到推理生态，性能数据详实，建议开发者立即评估迁移。
Google has open-sourced TorchTPU, a native PyTorch backend for Cloud TPUs. It was co-developed with the PyTorch team at Meta and announced at PyTorch Conference 2026 on October 20, 2026. Existing PyTorch code targets a TPU through a standard torch.device("tpu"). TPUs now sit beside GPUs and CPUs as a native PyTorch device. Google team states that TorchTPU replaces PyTorch/XLA and recommends it for all new workloads.
Google 已开源 TorchTPU，这是 Cloud TPU 的原生 PyTorch 后端。它由 Google 与 Meta 的 PyTorch 团队联合开发，并于 2026 年 10 月 20 日的 PyTorch Conference 2026 上宣布。现有的 PyTorch 代码通过标准的 torch.device("tpu") 针对 TPU 进行目标设定。现在，TPU 作为原生 PyTorch 设备，与 GPU 和 CPU 并列存在。Google 团队表示，TorchTPU 取代了 PyTorch/XLA，并推荐用于所有新工作负载。
TL;DR
TL;DR（太长不看）
- Size: A software stack, not a model. 4 open repos: TorchTPU, vLLM-TorchTPU, SGLang-TorchTPU, TorchTitan. TorchTPU is v0.2.0 on PyPI, Python 3.11+, Apache 2.0.
- Runs on: Google Cloud TPUs. Tested on generations from v5e, scaled to 8,192 Ironwood (TPU7x) chips. Serving supports FP8, MXFP4 and NVFP4 checkpoints.
- Performance: Eager mode for debugging, torch.compile for production-level throughput.
- Best: 52% MFU training Qwen 3 8B with TorchTitan on Ironwood (self-reported).
- Worst: Out-of-box eager took 11.49s per step in a fine-tuning case, versus 2.78s fully tuned.
- 规模：这是一个软件栈，而非模型。包含 4 个开源仓库：TorchTPU、vLLM-TorchTPU、SGLang-TorchTPU、TorchTitan。TorchTPU 在 PyPI 上的版本为 v0.2.0，支持 Python 3.11+，采用 Apache 2.0 许可证。
- 运行环境：Google Cloud TPUs。已在从 v5e 到 8,192 颗 Ironwood (TPU7x) 芯片的各代产品上经过测试。推理服务支持 FP8、MXFP4 和 NVFP4 检查点。
- 性能：Eager mode（急切模式）用于调试，torch.compile 用于生产级吞吐量。
- 最佳表现：在 Ironwood 上使用 TorchTitan 训练 Qwen 3 8B，MFU（模型浮点利用率）达到 52%（自我报告）。
- 最差表现：开箱即用的 eager 模式在一个微调案例中每步耗时 11.49 秒，而完全调优后为 2.78 秒。
What is TorchTPU?
什么是 TorchTPU？
TorchTPU is an open-source PyTorch backend from Google and Meta that runs PyTorch models natively on TPUs. It plugs in through PyTorch’s PrivateUse1 interface. Tensors are ordinary PyTorch tensors, with no subclasses or wrappers. It runs one process per device and preserves tensor views and in-place mutations.
TorchTPU 是 Google 和 Meta 推出的开源 PyTorch 后端，可在 TPUs 上原生运行 PyTorch 模型。它通过 PyTorch 的 PrivateUse1 接口接入。张量是普通的 PyTorch 张量，没有子类或包装器。它为每个设备运行一个进程，并保留张量视图和原地修改操作。
The predecessor, PyTorch/XLA, recorded lazy-tensor graphs. Developers often added code to mark compilable regions. PyTorch/XLA also only supported pure SPMD code. TorchTPU also handles divergent MPMD execution.
其前身 PyTorch/XLA 记录的是惰性张量图。开发者通常需要添加代码来标记可编译区域。PyTorch/XLA 也仅支持纯 SPMD 代码。TorchTPU 还能处理发散的 MPMD 执行。
How does TorchTPU execute PyTorch code?
TorchTPU 如何执行 PyTorch 代码？
TorchTPU ships 3 eager modes:
TorchTPU 提供三种 eager 模式：
- Strict (default): dispatches each operator asynchronously with standard PyTorch semantics.
- Debug: runs synchronously with tensor bounds checking.
- DeferAndFuse: fuses runs of eager operators to cut HBM traffic. The TorchTPU paper says it can be 50% or more faster than Strict.
- Strict（默认）：以标准 PyTorch 语义异步分发每个算子。
- Debug：同步运行并进行张量边界检查。
- DeferAndFuse：融合急切算子的运行以削减 HBM 流量。TorchTPU 论文指出，其速度可比 Strict 模式快 50% 以上。
For peak speed, torch.compile(backend="tpu") captures the graph with TorchDynamo. It lowers through StableHLO to OpenXLA, which emits TPU machine code. Compiled executables are cached across processes, hosts and job restarts.
为了获得峰值速度，torch.compile(backend="tpu") 使用 TorchDynamo 捕获计算图。它通过 StableHLO 降级到 OpenXLA，后者生成 TPU 机器码。编译后的可执行文件会在进程、主机和作业重启之间缓存。
Custom Pallas kernels are callable as standard PyTorch functions. Prebuilt Tokamax kernels and Helion tile kernels also work.
自定义 Pallas 内核可作为标准 PyTorch 函数调用。预构建的 Tokamax 内核和 Helion tile 内核也可用。
Distribution uses the tpu_dist process group. It supports DDP, FSDP2 and DTensor, along with a topology-aware mesh for TPU interconnects.
分布式使用 tpu_dist 进程组。它支持 DDP、FSDP2 和 DTensor，以及针对 TPU 互连拓扑感知的 mesh。
How fast is TorchTPU?
TorchTPU 有多快？
All figures below are self-reported by Google and partners.
以下所有数据均由 Google 及其合作伙伴自我报告。
- Training scale: TorchTitan reached up to 52% MFU on Qwen 3 8B. Google reports near-linear scaling up to 8,192 Ironwood TPUs.
- Tuning ladder: In a TorchTitan fine-tuning case, step time fell from 11.49s to 2.78s. DeferAndFuse alone gave 1.15x. Full torch.compile gave 3.80x.
- Model coverage: Google continuously tests about 5,000 Hugging Face models in eager and compiled modes.
- Fleet serving: Zenteiq grew daily output from 1.47B to 5.97B tokens per 1,000 TPU v6e chips. That gain came from llm-d routing on GKE, not TorchTPU alone.
- 训练规模：TorchTitan 在 Qwen 3 8B 上达到了高达 52% 的 MFU。Google 报告称，在多达 8,192 个 Ironwood TPU 上实现了近乎线性的扩展。
- 调优阶梯：在 TorchTitan 微调案例中，单步时间从 11.49 秒降至 2.78 秒。仅 DeferAndFuse 就带来了 1.15 倍的加速。完整的 torch.compile 带来了 3.80 倍的加速。
- 模型覆盖范围：Google 持续在 eager 模式和编译模式下测试约 5,000 个 Hugging Face 模型。
- 集群服务：Zenteiq 将每 1,000 个 TPU v6e 芯片的日输出量从 14.7 亿 token 提升至 59.7 亿 token。这一增益来自 GKE 上的 llm-d 路由，而非仅靠 TorchTPU。
What integrations ship with TorchTPU?
哪些集成随 TorchTPU 一起发布？
- Hugging Face: Transformers v5 deploys to TPU with device_map="tpu". Diffusers supports eager, torch.compile and tensor parallelism.
- vLLM-TorchTPU: built with Inferact. Keeps upstream schedulers and API servers, with custom Pallas kernels.
- SGLang-TorchTPU: built with RadixArk. Keeps RadixAttention prefix reuse.
- Serving features: FP8 KV caches, FP8, MXFP4 and NVFP4 checkpoints. Speculative decoding includes MTP and EAGLE-3.
- TorchTitan: recipes for Llama 3 70B/405B, Qwen 3 8B/32B, Qwen 3 235B MoE and FLUX.1 12B.
- Also: llm-d 1.0, Ray, TorchRec, VeRL (preview) and LiteRT-Torch for on-device export.
- Hugging Face：Transformers v5 使用 device_map="tpu" 部署到 TPU。Diffusers 支持 eager、torch.compile 和张量并行。
- vLLM-TorchTPU：与 Inferact 合作构建。保留了上游调度器和 API 服务器，并采用自定义 Pallas 内核。
- SGLang-TorchTPU：与 RadixArk 合作构建。保留了 RadixAttention 前缀复用。
- 服务特性：FP8 KV 缓存，FP8、MXFP4 和 NVFP4 检查点。投机解码包括 MTP 和 EAGLE-3。
- TorchTitan：提供 Llama 3 70B/405B、Qwen 3 8B/32B、Qwen 3 235B MoE 和 FLUX.1 12B 的配方。
- 还包括：llm-d 1.0、Ray、TorchRec、VeRL（预览版）以及用于设备端导出的 LiteRT-Torch。
How does TorchTPU compare with PyTorch/XLA, TorchNeuron and JAX?
TorchTPU 与 PyTorch/XLA、TorchNeuron 和 JAX 相比如何？
| Feature | TorchTPU (repo) | PyTorch/XLA (repo) | TorchNeuron (AWS docs) | JAX on TPU (Google docs) |
|---|---|---|---|---|
| Vendor | Google + Meta | Google / PyTorch | AWS | |
| Hardware | Cloud TPU (v5e and later tested) | Cloud TPU | Trainium 2, Trainium 3 | Cloud TPU |
| PyTorch device | tpu | xla | neuron | Not PyTorch |
| Backend mechanism | PrivateUse1 | Lazy-tensor graph recording | PrivateUse1 | Not applicable |
| Default execution | Eager (Strict, Debug, DeferAndFuse) | Lazy tensors | Eager plus torch.compile | Not disclosed |
| Compile path | TorchDynamo to StableHLO to XLA | XLA | Neuron Compiler | XLA |
| Distributed APIs | DDP, FSDP2, DTensor, SPMD and MPMD | Pure SPMD only | FSDP, DTensor, DDP, TP | Not disclosed |
| Custom kernels | Pallas, Tokamax, Helion | Not disclosed | NKI | Pallas |
| vLLM / SGLang serving | vLLM-TorchTPU, SGLang-TorchTPU | Not disclosed | vLLM Neuron | vLLM-JAX, SGLang-JAX |
| Status | Open source, v0.2.0 | Being replaced by TorchTPU | Public beta (Neuron 2.33) | Maintained |
| License | Apache 2.0 | Not disclosed | Not disclosed | Not disclosed |
| 特性 | TorchTPU (仓库) | PyTorch/XLA (仓库) | TorchNeuron (AWS 文档) | TPU 上的 JAX (Google 文档) |
|---|---|---|---|---|
| 厂商 | Google + Meta | Google / PyTorch | AWS | |
| 硬件 | Cloud TPU (已测试 v5e 及更高版本) | Cloud TPU | Trainium 2, Trainium 3 | Cloud TPU |
| PyTorch 设备 | tpu | xla | neuron | 非 PyTorch |
| 后端机制 | PrivateUse1 | 惰性张量图记录 | PrivateUse1 | 不适用 |
| 默认执行 | Eager (Strict, Debug, DeferAndFuse) | 惰性张量 | Eager 加上 torch.compile | 未公开 |
| 编译路径 | TorchDynamo 到 StableHLO 到 XLA | XLA | Neuron Compiler | XLA |
| 分布式 API | DDP, FSDP2, DTensor, SPMD 和 MPMD | 仅限纯 SPMD | FSDP, DTensor, DDP, TP | 未公开 |
| 自定义内核 | Pallas, Tokamax, Helion | 未公开 | NKI | Pallas |
| vLLM / SGLang 服务 | vLLM-TorchTPU, SGLang-TorchTPU | 未公开 | vLLM Neuron | vLLM-JAX, SGLang-JAX |
| 状态 | 开源，v0.2.0 | 正被 TorchTPU 取代 | 公测 (Neuron 2.33) | 维护中 |
| 许可证 | Apache 2.0 | 未公开 | 未公开 | 未公开 |
How do you install TorchTPU?
如何安装 TorchTPU？
Install the package on a TPU VM, then pick the tpu device:
在 TPU VM 上安装包，然后选择 tpu 设备：
# pip install torch_tpu
import torch
device = torch.device("tpu")
x = torch.randn((4, 4), device=device)The repo’s README pins torch_tpu==0.2.0 with the CPU PyTorch wheel index. Without it, pip pulls the CUDA build of torch. Full guides live at torchtpu.readthedocs.io.
该仓库的 README 将 torch_tpu==0.2.0 与 CPU PyTorch wheel 索引绑定。若无此配置，pip 会拉取 torch 的 CUDA 构建版本。完整指南位于 torchtpu.readthedocs.io。
Key Takeaways
关键要点
- TorchTPU is open source under Apache 2.0, co-built with Meta.
- It replaces PyTorch/XLA with eager-first, PrivateUse1-based execution.
- torch.compile(backend="tpu") lowers through StableHLO and XLA.
- TorchTitan hit 52% MFU on Qwen 3 8B, scaling to 8,192 Ironwood chips.
- vLLM, SGLang and Hugging Face now run on TPU through standard PyTorch.
- TorchTPU 在 Apache 2.0 许可下开源，由 Meta 联合构建。
- 它用基于 eager-first、PrivateUse1 的执行方式取代了 PyTorch/XLA。
- torch.compile(backend="tpu") 通过 StableHLO 和 XLA 进行编译。
- TorchTitan 在 Qwen 3 8B 上实现了 52% 的 MFU，并扩展至 8,192 个 Ironwood 芯片。
- vLLM、SGLang 和 Hugging Face 现在可通过标准 PyTorch 在 TPU 上运行。
Check out the GitHub repo, paper, docs and Google Cloud announcement. All credit goes to the researcher of this project. Also, feel free to follow us on Twitter and don’t forget to join our 150k+ML SubReddit and Subscribe to our Newsletter. Wait! are you on telegram? now you can join us on telegram as well.
请查看 GitHub 仓库、论文、文档以及 Google Cloud 公告。所有功劳归于本项目的研究人员。此外，欢迎在 Twitter 上关注我们，别忘了加入我们拥有超过 15 万成员的 ML SubReddit，并订阅我们的新闻通讯。等等！你在 Telegram 上吗？现在你也可以加入我们的 Telegram 群组。
更进一步：量化金融体系
看懂新闻只是起点——沿量化金融路径，把它变成能交付的工程能力