简介
TorchTPU 是 Google 在 2026 年 10 月 20 日 PyTorch Conference 2026 上宣布开源的 PyTorch 原生 TPU 后端，由 Google 与 Meta 的 PyTorch 团队共同开发，2025 年年中立项，已在部分大型 Google Cloud 客户的训练和推理生产环境中使用。与早期的 PyTorch/XLA 不同，它把 eager 执行放在首位，不依赖 JAX 运行时，现有 PyTorch 代码只需把设备从 cuda 改成 tpu；需要性能时再用 torch.compile 编译到 TPU。它支持 DDP、FSDP、DTensor 等分布式写法，并已接入 vLLM、SGLang、TorchTitan、Ray、TorchRec 与 Hugging Face Transformers / Diffusers，适合想在 TPU 上训练、微调或部署开源模型、又不想改用 JAX 的团队。安装方式为在 Google Cloud TPU VM 上 pip install torch_tpu，文档见 https://torchtpu.readthedocs.io/ 。项目仍在快速迭代，GitHub 页标注许可为 Other，README 写明 Apache 2.0。官方来源：https://cloud.google.com/blog/topics/ai-infrastructure/torchtpu-is-now-open-source-run-native-pytorch-on-tpus 、https://pytorch.org/blog/introducing-torchtpu/ 、https://docs.cloud.google.com/tpu/docs/pytorch-on-tpu 、https://github.com/google-pytorch/torch_tpu （2026-10-21 核对）。 定价：免费开源（Apache 2.0）；需在 Google Cloud TPU 上运行，TPU 算力按 Google Cloud 价格另计。
- PyTorch 原生 TPU 后端注册 "tpu" 设备，现有代码把 .to("cuda") 换成 .to("tpu") 即可运行，运行时不依赖 JAX
- 安装在 TPU VM 上执行 pip install torch_tpu（当前公开版本 0.2.0，Python 3.11+）
- eager 模式优先，便于逐步调试；新增 DeferAndFuse eager 融合；生产用 torch.compile（TorchDynamo → OpenXLA / StableHLO）编译
- 已支持约 450 个 ATen 算子；官方称在 Hugging Face 下载量前 5,000 的模型上验证可直接运行（Qwen、DeepSeek、Gemma、Llama、Mistral、Whisper、CLIP 等）
- 分布式DDP、FSDP、DTensor 等 PyTorch 分布式原语，集合通信走 torch.distributed（tpu_dist）
- 生态集成vLLM、SGLang、llm-d、TorchTitan、Ray、TorchRec、Hugging Face Transformers v5（device_map="tpu"）、Diffusers
- 自定义内核可用 JAX/Pallas 或 Helion DSL 编写
- 官方数据TorchTitan 训练 Qwen 3 8B 最高 52% MFU，在最多 8,192 颗 Ironwood（TPU v7x）上近线性扩展
- 文档与教程torchtpu.readthedocs.io（含 TPU VM / GKE 入门、FLUX 与 Muse Glimmer 30B 案例）
