Google推TorchTPU：PyTorch原生跑TPU，迁移只需改几行代码
Google 与 Meta 开源 TorchTPU，为 Google TPU 带来原生 PyTorch 支持。
Google从2020年起就借助PyTorch/XLA让PyTorch工作负载跑在TPU上，这条路子用到了JAX、XLA和延迟张量。按Google的说法，该方案短板相当明显。
具体问题有几处：PyTorch/XLA得对源码做侵入式改动，把可编译区域显式圈出来；编译延迟很严重；PyTorch原生的命令式“即时”执行模式它不支持；PyTorch生态里新出现的分布式计算原语，它的支持进度也偏慢。
TorchTPU换了一套思路，把即时执行放在最高优先级，现有PyTorch工作负载只要少量改动代码就能跑起来。
落到操作上，从GPU迁到TPU，需要把cuda 设备替换为 tpu，并将 nccl分布式后端替换成tpu_dist。此外开发者可以用新增的torch.accelerator.current_acceleratorAPI来识别当前可用的加速器，代码片段如下。
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提供Strict、Debug和DeferAndFuse三种模式。其中DeferAndFuse属于即时执行与编译执行混用的新玩法，Google称它“可以比 Strict Eager 快 50% 以上，同时降低 HBM 峰值内存占用”。
为检验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。
训练工作负载实验还用了Qwen3 MoE 235B，规模从1,024枚TPUv7芯片扩到8,192枚。推理这边，Google表示用vLLM跑大规模MoE工作负载，单位成本性能具备竞争力。
有开发者评价：“非常有意思的工作。Google 早就应该这样做了。PyTorch 的生态系统实在太丰富了。”
也有人提到，Google除通过GCP提供TPU外，还会把TPU卖给“用于本地部署的企业客户”，Anthropic就在其中，Meta未来也可能跟进。
TorchTPU眼下还在早期，仍在持续演进。路线图覆盖动态形状、编译性能、算子覆盖、强化学习和量化等方向。更多细节可查阅原始公告，其中对TorchTPU的架构和范围介绍得更全面。
