2026 年 10 月 21 日——Google 与 Meta 的 PyTorch 团队于 10 月 20 日在 PyTorch Conference 2026 公开 TorchTPU。项目以 Apache-2.0 协议放在 GitHub 仓库 google-pytorch/torch_tpu，PyPI 当前版本为 0.2.0，支持 Python 3.11 及以上。官方文档写明，TorchTPU 把 TPU 注册为与 GPU、CPU 并列的 PyTorch 设备后端，并取代 PyTorch/XLA；pytorch/xla 仓库页面标明将于 2026 年 11 月 15 日归档。PyTorch 官方博客在约 5000 个高下载 Hugging Face 模型中完成 4917 个端到端验证，其中 4821 个无崩溃通过，比例约 98%。
PyTorch 官方博客称，这项合作于 2025 年年中在 Google 启动。2020 年上线的 PyTorch/XLA 以 Lazy Tensor 接入 TPU，需要在源码中标出可编译区域，且没有原生 eager 执行。TorchTPU 走 PyTorch 调度器上的 C++ 后端，运行时不依赖 JAX，默认按一进程一设备执行，并保留张量 view 与原地修改语义。官方给出的单机改法是把 .to("cuda") 换成 .to("tpu")；分布式进程组后端从 nccl 改为 tpu_dist。同一篇博客写明，只改设备名可以让模型跑起来，峰值性能仍取决于形状、切分和算子是否适合 TPU 的脉动阵列。
调试、编译与安装
官方把 eager 分成三档。Strict 为默认，逐个 ATen 算子异步下发，对应 GPU 上的普通 eager。Debug 同步执行并打开边界检查，环境变量为 TORCH_TPU_LAUNCH_BLOCKING=1。DeferAndFuse 把多个算子推迟到物化点再融合编译，环境变量为 TORCH_TPU_DEFER_AND_FUSE=1。官方称该模式可比 Strict 快 50% 或更多，算子融合还会减少 HBM 上的临时张量。
生产路径使用 torch.compile(model, backend="tpu")。图由 TorchDynamo 捕获，降到 StableHLO，再由 XLA 生成 TPU 可执行文件。博客写明，当前在 TPU 上调用 torch.compile 会暂时固定 dynamic=False，因为有界动态形状尚未默认打开。编译缓存分三级：进程内内存、同机共享内存，以及可指向云存储的持久缓存；持久缓存默认关闭。
安装命令以仓库 README 为准。额外索引用于拉取 CPU 版 PyTorch，避免在 TPU 虚拟机上装入 CUDA 轮子：
pip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpu
import torch
device = torch.device("tpu")
x = torch.randn(4, 4, device=device)
y = torch.ones(4, 4, device=device)
print(torch.matmul(x, y).cpu().shape)
精度默认是 medium：矩阵乘输入用 bfloat16，累加用 FP32。high 与 highest 分别用 3 遍和 6 遍矩阵乘单元分解来逼近 FP32，官方给出的计算开销大约为 3 倍和 6 倍，并会触发性能警告。CUDA 上 PyTorch 的默认精度是 highest。博客称已实现约 450 个 ATen 算子，含主要架构所需 Core ATen 的前向与反向。自定义核可走 Pallas 接口 torch.tpu.pallas.jax_op（官方标为可用于生产），或实验性的 Helion DSL。
框架接入
Google Cloud 博客列出已接入的工具：Hugging Face、vLLM、SGLang、llm-d、TorchTitan、Ray、TorchRec、LiteRT。PyTorch 博客称，本次会议公开四个基础仓库：TorchTPU、vLLM-TorchTPU、SGLang-TorchTPU、TorchTitan。
Hugging Face 于 10 月 20 日发文：Transformers v5 用 device_map="tpu" 把模型直接加载到 TPU；Diffusers 支持 pipe.to("tpu")、torch.compile 和张量并行。该文写明，Transformers 与 Diffusers 可直接走 TorchTPU，Optimum TPU 退出这两条库的必需路径。Google 称已为 Transformers 套件建立 TPU 夜间 CI。
Google Cloud 文档写明，vLLM-TorchTPU 与 SGLang-TorchTPU 保留上游调度器、KV 缓存和 API 服务，计算密集块调用 Pallas 核。文档列举的开源模型包括 Qwen、DeepSeek、GLM、Kimi、GPT-OSS 和 Gemma 4。vLLM 仓库说明：TPU v7x（Ironwood）为推荐硬件，纳入每次拉取请求与夜间 CI；v6e、v5e、v5p 为实验支持，不在 CI 内。原先的 JAX 推理后端已更名为 vLLM JAX，文档称 vLLM-JAX 与 SGLang-JAX 继续维护。
TorchTitan 侧，Google Cloud 博客称 Qwen3 8B 训练最高达到 52% 模型浮点利用率（MFU），并在最多 8192 颗 Ironwood（TPU v7x）上做了近线性扩展验证。同一则公告还写入 VeRL 的 TPU 后端与 GRPO 参考配方，以及 Ray 对 CPU、GPU、TPU 的混部编排。
Meta PyTorch 工程总监 Jana van Greunen 在 Google Cloud 博客中表示，TorchTPU 用于把 PyTorch 负载部署到 Google TPU，覆盖主要的大语言模型、广告和推荐系统生产场景。PyTorch 博客的 TorchRec 案例写明，32 颗 TPU v7x 上的 DLRM v2，在使用基线 SparseCore 核、尚未加上进阶优化时，步时延为 18.3 毫秒，与 Google 的 JAX 基线相同。
官方报告的覆盖率与性能
下表来自 2026 年 10 月 20 日的 PyTorch 官方博客与 Google Cloud 博客。除注明转述外，均为项目方自测。
| 项目 | 数据 | 口径 |
|---|---|---|
| 模型开箱执行 | 4917 个中 4821 个无崩溃通过，约 98%；500 个以上架构族 | 高下载模型语料，eager 与编译模式，硬件自 TPU v5e 起 |
| 数值对齐 | 657 个代表架构中 638 个相对 GPU 通过，约 97% | 因果语言模型困惑度相对差不超过 0.10；视觉、音频、编码器要求 Relative MSE ≤ 0.10 或绝对 MSE < 0.05 |
| 有界动态形状 | 839 个因果语言模型中 834 个数值正确，约 99% | SymInt 专项评估；torch.compile 默认仍关闭动态形状 |
| Qwen3 8B 训练 | 最高 52% MFU | TorchTitan；Google Cloud 博客，最多 8192 颗 Ironwood |
| Qwen3 MoE 235B | 1024 颗 v7：每 TensorCore 2236 token/秒、MFU 28.8%、步时 29.3 秒；8192 颗：2104 token/秒、MFU 27.1%、步时 31.2 秒 | 序列长度 4096，每 TensorCore 批大小 16，BF16，SimpleFSDP / HSDP，含 Pallas 核 |
| 与 JAX 对照 | 64 颗 TPU v7：TorchTitan 1828 token/秒，MaxText 1937 token/秒 | 同一 Qwen3 MoE 235B 的冒烟对比 |
| Llama 3 405B | 1024 颗 TPU v7x、8K 上下文、BF16：每芯片 500 token/秒，MFU 超过 50% | TorchTitan |
| 纯 PyTorch 推理 | 单颗 TPU v7、序列长度 5120：约 2.57 万 token/秒，MFU 59% | Muse Glimmer 30B 预填充；一颗芯片对 PyTorch 暴露两个 TensorCore，示例使用 world size 为 2 的 FSDP |
| InferenceX | 4 颗 TPU v7 上的 Qwen 3.5，8K/1K：每百万 token 2.01 美分；博客称 9 月首次提交为 3.6 美分 | 博客转述 SemiAnalysis 评测 |
| AgentX | 每百万 token 0.83 美分 | 博客转述 Google 提交的首批结果，并写明所用为即将推出的 vLLM-TorchTPU 栈 |
| 推荐模型 | 32 颗 TPU v7x，DLRM v2 步时延 18.3 毫秒 | 基线 SparseCore 核，与 JAX 基线相同 |
Qwen3 MoE 235B 从 1024 颗芯片（总批大小 32768）扩到 8192 颗（总批大小 262144）时，每 TensorCore 吞吐从 2236 token/秒降到 2104 token/秒，MFU 从 28.8% 降到 27.1%。博客将这组结果描述为弱扩展：总吞吐随芯片数大致线性增长。多 slice 反向的 all-reduce 走数据中心网络，由 XLA 与相邻层计算重叠。
上表性能是项目方在指定模型、并行策略和核实现上的自测，InferenceX 数字由 PyTorch 博客转述 SemiAnalysis。开箱通过率衡量的是不崩溃，数值对齐只覆盖 657 个代表模型。默认
medium精度使用 bfloat16 矩阵乘，与 CUDA 上默认的 FP32 路径不同。官方 FAQ 写明 TorchTPU 与 JAX 并存，JAX 训练和服务栈继续可用。新的 PyTorch TPU 任务按官方建议从 TorchTPU 开始；已有 PyTorch/XLA 负载迁移到 TorchTPU，该仓库页面写明的归档日期为 2026 年 11 月 15 日。
举手提问