Google 让 PyTorch 原生跑上 TPU：5000 个模型验证，迁移只需改几行

Google 从 2020 年起就借助 PyTorch/XLA 让 PyTorch 工作负载能跑在 TPU 上，这套方案依赖 JAX、XLA 和延迟张量。Google 自己也承认，这种做法存在明显短板。

PyTorch/XLA 的问题出在几个方面：开发者必须对源代码做侵入式修改，才能显式划出可编译区域；编译延迟相当严重；PyTorch 原生的命令式（即时）执行模式得不到支持；对 PyTorch 生态里新冒出来的分布式计算原语，支持进度也一直落后。

TorchTPU 换了一条路子，把即时执行放在最高优先级，现有 PyTorch 工作负载只需少量代码修改就能跑起来。具体来说，从 GPU 迁移到 TPU，需要将

分布式后端替换为cuda 设备替换为 tpu，并将 nccl。开发者也可以使用新的tpu_distAPI 识别可用的加速器，如以下代码片段所示。torch.accelerator.current_acceleratorTorchTPU 提供三种即时执行模式：Strict、Debug 和 DeferAndFuse。其中 DeferAndFuse 是把即时执行与编译执行混在一起的新模式，Google 称它“可以比 Strict Eager 快 50% 以上，同时降低 HBM 峰值内存占用”。

import torchimport torch.distributed as distdef 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 能否在不改动代码的前提下运行现有 PyTorch 模型，Google 搭了一条集成流水线，在从 v5e 开始的多代 TPU 硬件上，用即时和编译两种模式评估了约 5,000 个 Hugging Face 模型，覆盖 500 多种架构。

测试方法上，Google 采用了一套基于语料库分层和架构均衡选择的稳健测试方法。

这套方法走两级选择流程，在现实采用情况与广泛的架构多样性之间取得平衡：既避免重复测试细微变体，又确保真实生产权重得到充分评估。

性能验证方面，Google 跑了一个用标准 PyTorch 实现、拥有 300 亿参数的 Muse Glimmer 模型。单枚 TPUv7 芯片上，序列长度 5120 时推理吞吐量达到每秒 25,700 个 Token，序列长度 10,240 时为每秒 22,400 个 Token。训练侧，使用四枚 TPUv7 芯片，序列长度 5120 和 10,240 对应的吞吐量分别是每秒 5,600 和 4,600 个 Token。

Google 还用 Qwen3 MoE 235B 做了训练工作负载实验，规模从 1,024 枚 TPUv7 芯片扩到 8,192 枚。推理方面，Google 表示用 vLLM 运行大规模 MoE 工作负载，单位成本性能具有竞争力。

有开发者评价：“非常有意思的工作。Google 早就应该这样做了。PyTorch 的生态系统实在太丰富了。”

其他用户指出，Google 除了通过 GCP 提供 TPU，还会将其出售给“用于本地部署的企业客户”，其中包括 Anthropic，而 Meta 未来也可能跟进。

TorchTPU 目前还很早期，功能也在不断演进。动态形状、编译性能、算子覆盖、强化学习、量化，都是这个项目路线图里列出的方向。想了解 TorchTPU 的架构和范围，可以看原始公告，里面有更完整的说明。