谷歌开源 TorchTPU，PyTorch 可原生运行于 TPU

编辑｜Panda 如果只能用一行代码概括这条新闻，那就是 device_map="tpu"。 在圣何塞举行的 PyTorch Conference 2026 上，谷歌宣布开源 TorchTPU，这是一套让 PyTorch 原生运行在 TPU 上的完整软件栈。从此，在 Hugging Face Transformers v5 里把模型部署到 TPU，只需要传入这一个参数。 一个参数的变化，背后是谷歌在软件路线上的一次转向，也是英伟达 CUDA 生态近十年来遭遇的最正面的一次挑战。 AI 芯片的竞争，从来不只是芯片本身的竞争。PyTorch 是全球使用最广泛的 AI 框架，而它自 2016 年诞生起就与 CUDA 深度绑定：英伟达工程师长期投入，确保 PyTorch 在 N 卡上跑得最好，华尔街甚至有分析师把 CUDA 视为英伟达抵御对手最坚固的盾牌。反观谷歌，TPU 十年来主要服务自家的 TensorFlow 和 JAX。PyTorch 用户想用 TPU，得经过 PyTorch/XLA 这样一层「转接」，付出大量额外的工程成本。 TorchTPU 改变的正是这个前提。用谷歌官方的说法，TPU 从此与 GPU、CPU 并列，成为 PyTorch 的原生选项。 项目地址：https://github.com/google-pytorch/torch_tpu 项目由谷歌与 PyTorch 的「娘家」Meta 联合开发，vLLM、SGLang、Hugging Face、Ray、TorchTitan 等主流开源项目同步完成接入，一份 23 页的技术报告也同时发布在 arXiv 上。 报告标题：TorchTPU – A Hybrid Eager/Compiled Architecture for Native PyTorch Execution on TPU 报告地址：https://arxiv.org/pdf/2610.24161 时间点也值得关注。过去一年，Anthropic 签下最多 100 万颗 TPU 的使用协议，Meta 也被曝与谷歌达成数十亿美元的 TPU 租用合约。硬件已经卖得出去，软件成了最后一块短板。现在，谷歌要亲手把它补上。 不是又一个「转接头」 PyTorch 跑 TPU 并不新鲜。技术报告回顾，谷歌早在 2020 年就推出了基于惰性张量的 PyTorch/XLA，并罕见地直言了它的问题：需要侵入式修改源码来划定可编译区域，编译延迟严重，不支持 PyTorch 原生的即时执行模式，对分布式新原语的跟进也很滞后。后来的 TorchAX 则把 PyTorch 算子逐个翻译成 JAX 执行。在报告看来，这些方案本质上都是让 PyTorch 去迁就 JAX 的运行方式。 TorchTPU 在 2025 年年中立项时就定下几条原则：运行时不依赖 JAX，即时执行必须扎实，把代码从 GPU 迁到 TPU 应当像把 .to("cuda") 改成 .to("tpu") 一样简单。 实现上，它基于 PyTorch 预留给第三方硬件的 PrivateUse1 接口，在 PyTorch 分发器中注册为一个新的 C++ 后端，开发者拿到的就是普通的 PyTorch 张量，没有子类，也没有包装器。 从 GPU 迁移到 TPU 前后的设备初始化代码对比 即时执行被做成三档： 默认的 Strict 模式逐个分发算子、与 CPU 异步执行，对应 GPU 上的常规体验； Debug 模式同步执行，用于排查形状错误和 NaN； 第三档 DeferAndFuse 是 TorchTPU 新引入的即时与编译混合模式，先把一串算子攒起来，到需要取值时再自动融合成大块计算交给 TPU，速度比 Strict 模式快 50% 以上，还能降低峰值显存占用。 三种模式共享一套分层编译缓存，同样的算子序列编译一次即可反复复用。 追求极致性能时，开发者可以切换到 torch.compile。这里谷歌做了一个刻意的取舍：用 TorchDynamo 捕获计算图后，不走 PyTorch 默认的 Inductor 编译器，而是把算子映射到 StableHLO，交给 XLA 生成 TPU 程序。 报告给出了理由：一是能复用为即时模式搭好的基础设施，二是 XLA 经过多年检验足够稳定，三是 Inductor 的成功很大程度上建立在 Triton 之上，而 Triton 很难移植到 TPU 架构。这句话从侧面说明，PyTorch 与 GPU 的绑定，深入到了编译器这一层。 在 TPU 上调用 torch.compile 只需指定 backend=“tpu” PyTorch 经由 StableHLO 接入 OpenXLA 编译栈，与 JAX、Pallas 内核共用同一条下降路径 分布式方面

来源只提供摘要，完整内容请打开原文。