PyTorch原生支持TPU，谷歌与Meta联合开源TorchTPU，模型迁移仅需一行代码

在圣何塞举行的 2026 年 PyTorch 北美大会主题演讲中，谷歌与 Meta 联合宣布开源 TorchTPU——一个深度集成进 PyTorch 核心调度器的原生 TPU 后端。这一举措被视作 PyTorch 生态向多硬件统一支持迈出的关键一步。

TorchTPU 的核心目标是消除模型在不同硬件平台间迁移的技术门槛。此前，PyTorch 用户若想在谷歌 TPU 上运行模型，需依赖 PyTorch/XLA 这一桥接层。该方案自 2020 年推出以来，虽实现基本功能，但要求开发者对原始代码进行侵入性修改，例如重写张量操作逻辑以适配 XLA 的惰性执行模型。Meta 软件工程师 Zongwei Zhou 在演讲中明确指出，新后端的设计哲学是“让迁移像将 model.to('cuda') 改为 model.to('tpu') 一样简单”。

该后端通过 PyTorch 的 PrivateUse1 扩展机制注册，与近期亚马逊云科技推出的 TorchNeuron（用于 Trainium 芯片）采用相同技术路径。这意味着未来更多定制化 AI 加速芯片有望以标准化方式接入 PyTorch 生态，减少厂商锁定风险。

据谷歌云 Core ML/AI 工程副总裁 Bill Jia 介绍，TorchTPU 已实现与主流推理与训练框架的兼容：Hugging Face 模型无需修改即可加载；vLLM 和 SGLang 推理引擎可原生运行；TorchTitan 支持密集型、专家混合型（MoE）及扩散模型的大规模训练；而 Ray 与 llm-d 则可在 Kubernetes 集群中跨多台主机调度 TPU 工作负载。目前已开源的四个代码库包括：TorchTPU 本体、vLLM-TorchTPU、SGLang-TorchTPU 与 TorchTitan。

性能方面，Jia 展示了在 8192 个 TPU v7x 芯片组成的超大规模集群上运行 Qwen3 235B MoE 模型的结果：模型 FLOPs 利用率达 28%，单芯片吞吐量在集群扩展至数千芯片时仍保持近线性增长。这表明 TPU 架构在超大规模分布式训练场景下具备良好的可扩展性。

不过，公告也坦承：仅靠一行 .to('tpu') 修改尚不足以发挥 TPU 的峰值性能。由于模型结构、张量分片策略与算子实现往往针对 GPU 架构优化，直接迁移可能带来性能折损。对此，谷歌提出的解决方案并非强制重构模型，而是引入“TPU 代理套件”（TPU Agent Suite），该工具已在 GitHub 开放预览，可自动完成环境配置、故障诊断与性能调优，降低用户使用门槛。

这一进展背后，是 AI 工作负载结构的深刻变化。PyTorch 基金会执行董事 Mark Collier 指出，推理已取代训练成为数据中心最大负载，“对 token 的需求”急剧上升。Agentic AI 基金会执行董事 Mazin Gilbert 提供的数据进一步印证：2026 年 2 月 6 日起，智能体产生的 token 消耗量已超过人类总消耗，并在后续数月内增长 14 倍，远高于人类同期 2.8 倍的增长。这种“曲棍球杆式”爆发曲线，倒逼基础设施必须具备更高弹性与跨平台兼容能力。

值得注意的是，TorchTPU 的推出并非孤立事件，而是行业协同演进的一部分。Meta 正在加码 AI 基础设施投入，计划募资 950 亿美元用于建设超级计算集群；SpaceX 与阿波罗全球管理签署约 400 亿美元 GPU 主租赁协议；英伟达与甲骨文合作部署 Vera Rubin NVL72 超级集群；AWS 与 Verizon 签署超 20 亿美元暗光纤协议——所有这些动作，均指向同一方向：构建高带宽、低延迟、异构兼容的 AI 算力底座。

硬件层面，英伟达 Vera CPU 的实测数据亦揭示新趋势：智能体工作负载高度依赖 CPU 执行环节（如搜索、编码、测试），其瓶颈已从纯 GPU 算力转向“CPU-GPU 协同效率”与“每瓦吞吐量”。Vera CPU 在 100 兆瓦数据中心场景下，相较最优 x86 处理器，每瓦可支持 1.6 倍智能体数量，速度提升达 2 倍。Perplexity、DeepInfra、ClickHouse 等早期合作方已验证其在真实任务中的优势。

回到 TorchTPU 本身，其意义不仅在于技术实现，更在于生态共识的建立。Meta PyTorch 工程总监 Jana van Greunen 强调：“推动前沿 AI 发展需要巨大的计算规模和丰富的硬件多样性，我们不能被单一架构束缚。”谷歌与 Meta 的联合行动，实质是将“硬件无关性”从理念推向工程现实。

目前，TorchTPU 已全部开源，开发者可立即尝试。尽管初期性能调优仍需人工介入，但随着代理套件与社区优化持续迭代，TPU 有望真正成为 PyTorch 的“第一类原生目标”。在智能体驱动的 AI 新周期中，工作负载的可移植性，正成为比单点算力更重要的基础设施能力。