谷歌Meta联手推出TorchTPU：PyTorch原生拥抱TPU

谷歌与Meta近日在PyTorch官方博客联合宣布推出TorchTPU，这是一个开源协作项目，为谷歌TPU提供原生、规范化的PyTorch后端。该项目于2025年年中在谷歌内部启动，随后与Meta的PyTorch团队展开合作，目标是让PyTorch开发者无需大改代码即可在TPU上运行模型，并通过eager执行优化、torch.compile集成和原生硬件加速，提供可与GPU竞争的推理与训练性能。

从PyTorch/XLA到TorchTPU：一次架构理念的转变

TPU历来与TensorFlow和JAX生态深度绑定。2020年，谷歌曾推出首个面向TPU的编译版PyTorch——基于Lazy Tensors的PyTorch/XLA。这一方案虽然具有奠基意义，但难以跟上社区驱动的PyTorch快速演进：它需要对源代码进行侵入式修改以显式定义可编译区域，编译延迟严重，不支持PyTorch原生的命令式eager执行模式，在分布式计算原语的支持上也明显滞后。

与PyTorch/XLA和TorchAX等把PyTorch硬塞进JAX工作方式的前作不同，TorchTPU被构建为PyTorch的原则化TPU后端。团队制定了一套明确的设计原则：运行时不依赖JAX；提供扎实的eager模式；torch.compile基于Torch Dynamo（而非Lazy Tensor）实现，性能对标JAX编译；通过在PyTorch栈最底层（如Process Group API）提供TPU集合通信，广泛支持DDP、FSDP、DTensor等分布式习惯用法；并无缝融入PyTorch生态，积极 engages 社区。

以用户为中心：两行代码迁移TPU

TorchTPU的核心目标是可移植性——把工作负载迁到TPU应像把model.to("cuda")改成model.to("tpu")一样简单。官方坦言，仅靠这一改动无法达到TPU峰值性能，因为模型的形状、分片和算子可能偏向GPU架构，但模型应当开箱即跑。得益于PyTorch的autoload机制，在安装了TorchTPU wheel的TPU机器上，导入PyTorch后只需几行代码即可完成切换，分布式后端从"nccl"改为"tpu_dist"即可，也可以使用与GPU/TPU无关的current_accelerator写法。

团队还强调提供用户友好、可操作的运行时与编译器错误信息，并集成PyTorch社区已在使用的工具与框架，如torch.profiler、DTensor、Ray等。

TPU硬件与编程模型：MXU、SparseCore与XLA栈

TPU是专为深度学习张量运算设计的专用ASIC。从硬件架构看，一颗TPU芯片（依代际不同）包含一个或多个TensorCore，每个TensorCore由一个或多个矩阵乘法单元（MXU）、向量单元（VPU）和标量单元组成；MXU可包含256×256乘累加器，以脉动阵列方式排列。芯片还可能配备SparseCore，专门处理稀疏操作（如gather、scatter、embedding）和集合通信，可与TensorCore并行运行，实现计算与通信重叠。

内存与互联方面，TPU 8i拥有288 GB高带宽内存（HBM），带宽达8.6 TB/s；TPU 8t与8i的双向芯片间互联（ICI）带宽为19.2 TBps，支撑用户配置的TPU Slice（TPU Pod的逻辑子组，以3D环面或TPU 8i的Boardfly拓扑连接）；TPU 8t的双向数据中心网络（DCN）带宽为400 GBps，用于跨Slice通信。TPU通过XLA编译器栈编程，入口为StableHLO（较高层ML IR）和Mosaic（较低层自定义kernel IR），TPU实际ISA专有且不公开，StableHLO尤其是Mosaic可类比为TPU版的PTX。C++程序通过PjRT Client访问本地TPU，编译MLIR为可执行文件并异步启动。

编程模型上有几个值得注意的特点：TensorCore可视为单线程处理器，模型形状需足够宽才能榨干MXU算力，这与GPU可并行启动多个小kernel的思路不同；分布式方面TorchTPU采用与JAX一进程一主机不同的"一设备一进程"模型，但由于一颗TPU只能被一个用户态进程访问，RL colocation暂时无法开箱支持，团队正在研究绕过方案；集合通信算子直接lower为StableHLO算子，而非使用NCCL这样的独立库；此外团队正在引入有界动态形状（bounded dynamism），在编译期已知维度上限、运行期才确定实际维度时避免XLA重复编译。

数值精度策略与性能权衡

TorchTPU将PyTorch的标准精度API适配到TPU硬件上。TPU原生矩阵乘法格式为bfloat16（BF16），MXU以BF16输入、32位浮点中间累加的方式执行运算，这是TPU高性能计算的基石。对于非收缩类操作（逐元素加法、激活函数、LayerNorm、Softmax等），TorchTPU在VPU上执行完整IEEE-754 FP32运算。

TorchTPU提供三档精度模式：medium（默认）对乘法使用BF16、累加使用FP32，足以满足几乎所有深度学习工作负载；high采用三趟算法，将输入拆分为高低两部分计算三个交叉项，以约3倍计算代价恢复大部分FP32精度；highest采用六趟算法逼近完整FP32精度，代价约6倍。官方建议绝大多数场景使用默认medium档，显式选择high或highest会触发TORCH_WARN_ONCE性能警告，只有在确认出现精度相关的收敛问题时才应切换。

当前进展：450个算子与5000个模型验证

TorchTPU仍在积极开发中，但已被谷歌云部分最大客户用于推理与训练的生产环境，代码已在GitHub开源。截至发稿，项目已支持eager与编译两种模式，包括新增的DeferAndFuse eager模式；支持约450个Aten算子（跳过具有结构性或复合委托的算子），覆盖主流深度学习架构所需全部Core ATen算子的前向与反向传播；并支持通过JAX/Pallas和Helion DSL编写自定义kernel。

兼容性验证成果尤为亮眼：团队在最热门的5000个Hugging Face模型上验证了开箱即用的执行，包括LLM领域的Qwen、DeepSeek、Gemma、Llama、Mistral、GPT、OPT、BLOOM、Falcon和StableML；语音领域的Whisper、Wav2Vec2、HuBERT和MMS；视觉领域的DinoV2、SigLIP、CLIP、Swin Transformer、ConvNeXt、ViT、ResNet和DETR；文本编码器与搜索领域的BERT、ALBERT、Electra、RoBERTa、DeBERTa、T5、MT5和Flan-T5。

此外，项目还完成了基于DTensor的分片验证、多主机推理（vLLM和SGLang）与微调/训练（TorchTitan，配合Ray编排）验证，以及与vLLM、SGLang、TorchTitan、Ray和TorchRec的集成验证。

影响与展望：PyTorch生态的第三条硬件路径

在AI算力竞争白热化的当下，TorchTPU的意义远超一次技术发布。对PyTorch数百万开发者而言，这意味着除CUDA GPU外又多了一条经过谷歌与Meta两大巨头背书的硬件路径；对谷歌而言，降低TPU的迁移门槛将直接扩大其自研芯片在训练与推理市场的吸引力，与英伟达形成更直接的竞争。由于模型形状、分片策略仍需针对TPU架构调优才能达到峰值性能，短期内迁移成本依然存在，TorchTPU也尚未解决RL colocation等场景的限制。但凭借生产环境验证、开源协作模式以及对PyTorch原生习惯用法的坚持，TorchTPU有望推动"一次编写、多硬件运行"的愿景向前迈出实质性一步，未来PyTorch在GPU与TPU之间的大规模工作负载竞争值得持续关注。