Google发布TorchTPU让PyTorch原生支持TPU

Google与Meta合作推出开源项目TorchTPU，为Google TPU提供原生PyTorch后端。现有PyTorch工作负载经过少量代码调整后，即可运行在TPU上，并获得即时执行、编译执行、分布式计算和硬件优化能力。

此前Google自2020年起主要通过PyTorch/XLA支持TPU，该方案依赖JAX、XLA和延迟张量。Google认为，PyTorch/XLA需要显式划定可编译区域，代码侵入较大，编译等待时间长，也无法原生支持PyTorch的命令式即时执行模式，对新分布式计算原语的支持也相对滞后。TorchTPU则将即时执行置于优先位置，降低迁移成本。

从GPU迁移时，开发者主要需要把cuda设备替换为tpu，并将nccl分布式后端改为tpu_dist；也可以调用torch.accelerator.current_accelerator统一识别当前加速器。TorchTPU提供Strict、Debug和DeferAndFuse三种即时执行模式，其中DeferAndFuse结合即时与编译执行，Google称其速度相比Strict Eager可提升超过50%，同时能够降低HBM峰值占用。

为验证兼容性，Google搭建了覆盖多代TPU的测试流水线，从v5e开始，在即时和编译两种模式下评估约5000个Hugging Face模型，涉及500多种架构。测试采用基于语料库分层和架构均衡的两级筛选方法，减少细微变体的重复测试，同时覆盖真实生产权重和多样化模型结构。

性能测试中，标准PyTorch实现的300亿参数Muse Glimmer模型在单枚TPUv7芯片、序列长度5120和10240时，推理吞吐量分别达到每秒25700和22400个Token；使用四枚TPUv7芯片训练时，对应吞吐量为每秒5600和4600个Token。另一项实验将Qwen3 MoE 235B训练规模从1024枚TPUv7扩展至8192枚，Google还表示，vLLM运行大规模MoE推理时可实现有竞争力的单位成本性能。

TorchTPU目前仍处于早期阶段，后续计划覆盖动态形状、编译性能、算子覆盖、强化学习和量化等方向。社区评论认为，PyTorch生态丰富，原生TPU支持早就应当推出；也有讨论提到，除通过云服务提供TPU外，Google可能向包括Anthropic在内的企业客户销售本地部署设备，Meta未来或许也会跟进。