11月12日,Google与Meta合作开发的开源项目TorchTPU亮相,为Google TPU提供原生PyTorch后端,使PyTorch工作负载只需少量代码修改即可在TPU上执行,并支持即时执行、编译执行、分布式计算及硬件专属优化。

Google自2020年起通过PyTorch/XLA支持PyTorch在TPU上运行,但该方案需侵入式修改源码以显式定义可编译区域,编译延迟严重,不支持PyTorch原生命令式执行,且在支持新兴分布式原语方面进展滞后。TorchTPU将即时执行列为最高优先级,从GPU迁移到TPU只需将cuda设备替换为tpu、将nccl后端替换为tpu_dist,开发者也可通过torch.accelerator.current_accelerator API识别可用加速器。

TorchTPU支持Strict、Debug和DeferAndFuse三种即时执行模式,其中DeferAndFuse混合模式据称可比Strict Eager快50%以上,同时降低HBM峰值内存占用。Google构建集成流水线,在从v5e开始的多代TPU上以即时和编译两种模式评估约5000个Hugging Face模型,覆盖500多种架构。

性能验证方面,300亿参数的Muse Glimmer模型在单枚TPUv7芯片上,序列长度5120时推理吞吐达每秒25700个Token,序列长度10240时为每秒22400个Token;使用四枚TPUv7训练时,两种序列长度吞吐分别为每秒5600和4600个Token。Qwen3 MoE 235B训练实验则将规模从1024枚TPUv7扩展至8192枚。TorchTPU仍处早期阶段,路线图涵盖动态形状、编译性能、算子覆盖、强化学习和量化等方向。