Google 与 Meta 联合推出 TorchTPU，PyTorch 原生支持 TPU 硬件

背景

Google 从 2020 年开始通过 PyTorch/XLA 让 PyTorch 跑在 TPU 上，但这套方案问题明显。开发者必须在源码里显式标记可编译区域，编译延迟严重，不支持 PyTorch 原生的即时执行模式，对新兴的分布式计算原语支持滞后。这些限制让许多 PyTorch 用户望而却步。

TorchTPU 是 Google 和 Meta 合作开发的新方案，把即时执行放在首位。从 GPU 迁移到 TPU 只需要两处改动：把设备名从 cuda 改成 tpu，把分布式后端从 nccl 改成 tpu_dist。如果用新的 torch.accelerator.current_accelerator API，连设备名都不用手动指定，直接识别可用加速器。

python 复制代码

import torch
import torch.distributed as dist

def main():
    device = torch.device("cuda") # -> device = torch.accelerator.current_accelerator
    dist.init_process_group(backend="nccl") # -> dist.init_process_group()
    rank = dist.get_rank()
    word_size = dist.get_world_size()
    ...
    model = model.to(device)

核心原理

TorchTPU 提供三种即时执行模式：Strict、Debug 和 DeferAndFuse。DeferAndFuse 是一种混合模式，兼顾即时执行和编译执行的优点，Google 称它比 Strict 模式快 50% 以上，峰值内存占用更低。

为了验证兼容性，Google 搭建了一条集成流水线，在从 v5e 开始的多代 TPU 硬件上测试约 5000 个 Hugging Face 模型，覆盖 500 多种架构。测试分两级选择，既平衡真实采用情况，又保证架构多样性，避免重复测试细微变体。

性能测试用了两个模型。一个是 300 亿参数的 Muse Glimmer，标准 PyTorch 实现。单枚 TPUv7 芯片上，序列长度 5120 时推理吞吐量达到每秒 25700 个 Token，序列长度 10240 时达到每秒 22400 个 Token。训练用四枚 TPUv7 芯片，两种序列长度下吞吐量分别是每秒 5600 和 4600 个 Token。

另一个是 Qwen3 MoE 235B，训练从 1024 枚 TPUv7 芯片扩展到 8192 枚。推理用 vLLM 跑大规模 MoE 工作负载，Google 表示单位成本性能有竞争力。

局限

TorchTPU 还处于早期阶段。路线图上列了动态形状、编译性能、算子覆盖、强化学习和量化等方向。这意味着部分功能还不完善，可能遇到未覆盖的算子或性能瓶颈。

项目依赖 TPU 硬件，测试覆盖从 v5e 开始的几代产品。旧硬件支持情况不明确。迁移代码虽然简单，但分布式场景下调试复杂度可能比文档展示的高。

Google 在公告里提到企业客户可以购买 TPU 用于本地部署，Anthropic 已经在用，Meta 未来可能跟进。但对于只能通过 GCP 使用 TPU 的团队，成本和可用性仍然是考虑因素。

对开发者的意义

PyTorch 生态足够丰富，让 PyTorch 原生跑在 TPU 上是不少人期待已久的事。Hacker News 上有用户直言"Google 早就应该这样做了"。TorchTPU 降低了迁移门槛，现有 PyTorch 代码几乎可以零修改运行，这对已经在用 PyTorch 的团队有吸引力。

三种即时执行模式给了不同场景的选择空间。DeferAndFuse 模式在性能和内存占用上的改进，对训练大模型的场景有实际意义。

测试覆盖 5000 多个模型和 500 多种架构，说明兼容性不是空话。但项目仍在演进，生产环境使用前需要充分评估自己的模型和工作负载是否在覆盖范围内。

TPU 的定价和 GPU 不同，单位成本性能可能更优，但需要结合具体工作负载计算。对于已经在 GCP 上运行 PyTorch 的团队，TorchTPU 是一个值得尝试的选项。对于本地部署或混合云场景，需要评估 TPU 硬件采购和维护成本。