谷歌开源TorchTPU，PyTorch第一次把TPU当成原生设备，CUDA最深的护城河松了
device_map="tpu"：谷歌开源 TorchTPU，PyTorch 第一次把 TPU 当成原生设备，CUDA 最深的护城河松了
Part One：一个参数，为什么值得专门写一篇
2026 年 10 月的 PyTorch Conference 2026（圣何塞）上，谷歌宣布开源 TorchTPU——一套让 PyTorch 原生运行在 TPU 上的完整软件栈，由谷歌与 PyTorch 的"娘家"Meta 联合开发，以 Apache 2.0 许可放出。同步动作很整齐：Hugging Face Transformers v5 里把模型部署到 TPU，只需传一个参数；vLLM、SGLang、Ray、TorchTitan 等主流项目完成接入；一份 23 页技术报告挂上 arXiv（2610.24161，题为《TorchTPU – A Hybrid Eager/Compiled Architecture for Native PyTorch Execution on TPU》）。
如果只能用一行代码概括这条新闻，那就是：
# 以前:PyTorch 想上 TPU,得经过 PyTorch/XLA 这层"转接头"
# 现在:TPU 与 GPU、CPU 并列,成为 PyTorch 的原生设备
model.to("tpu") # 就像当年写 .to("cuda") 一样自然
# 或者在 transformers v5 里:
AutoModelForCausalLM.from_pretrained(name, device_map="tpu")
为什么一个"加了个设备名"的更新值得写？因为 AI 芯片的竞争，从来不只是芯片本身的竞争。过去十年，英伟达最深的护城河不是晶体管，而是 CUDA 生态——PyTorch 自 2016 年诞生起就与 CUDA 深度绑定，英伟达工程师长期投入保证它在 N 卡上跑得最好，华尔街分析师直接把 CUDA 称作英伟达抵御对手最坚固的盾牌。而 TPU 十年来主要服务谷歌自家的 TensorFlow 和 JAX，PyTorch 用户想碰 TPU，必须绕道 PyTorch/XLA，付出真实的工程代价。
TorchTPU 改变的正是这个前提。谷歌官方的说法是：TPU 从此与 GPU、CPU 并列，成为 PyTorch 的原生选项。 时机也耐人寻味——过去一年 Anthropic 签下最多 100 万颗 TPU 的使用协议、Meta 被曝与谷歌达成数十亿美元 TPU 租用合约。硬件已经卖得出去，软件成了最后一块短板，谷歌现在亲手把它补上。
Part Two：老路子为什么痛——XLA 是让 PyTorch 迁就 JAX
要理解 TorchTPU 的不同，先看此前 PyTorch 用户上 TPU 的唯一主路：PyTorch/XLA。
XLA（Accelerated Linear Algebra）本是为 TensorFlow/JAX 设计的图编译器。PyTorch/XLA 做的事，本质是在 PyTorch 和 XLA 之间架一座桥：把 PyTorch 的动态算子调用"录制"成一张计算图（XLA graph），编译后再下发给 TPU。这套方案能用，但有几个结构性的别扭：
PyTorch/XLA(旧:转接头) TorchTPU(新:原生后端)
+-----------------------------+ +------------------------------+
| PyTorch eager 代码 | | PyTorch eager 代码 |
| ↓ 录制成图 | | ↓ 直接分发 |
| XLA graph(迁就JAX的运行方式) | => | PyTorch 分发器(dispatcher) |
| ↓ 编译 | | ↓ 注册的 C++ 后端(PrivateUse1) |
| TPU | | TPU(混合 eager/compiled 运行时) |
+-----------------------------+ +------------------------------+
问题:图边界/重编译/动态形状卡顿 原则:运行时不依赖JAX,eager要扎实
第一，运行模型被 XLA 的编译模型绑架。PyTorch 的灵魂是 eager（即时执行，逐行立即出结果，方便调试），而 XLA 倾向于把一大段图编译后整体执行。两者哲学相冲突，开发者经常要和"图边界（graph boundary）““懒惰张量（lazy tensor）““意外重编译（recompilation）“搏斗。
第二，动态形状和动态控制流容易触发重编译。Agent、变长序列这类 workload 形状一直在变，每变一次可能就重新编译一次，首步延迟（compile latency）很难看。
第三，调试体验割裂。报错常常来自底层编译后的图，而不是你写的那行 Python，栈追踪对不上。
报告的判断很直接：这些方案本质上都是"让 PyTorch 去迁就 JAX 的运行方式”。TorchTPU 在 2025 年年中立项时就定下几条硬原则：运行时不依赖 JAX；即时执行必须扎实（不能只是个样子）；把代码从 GPU 迁到 TPU 应当像把 .to("cuda") 改成 .to("tpu") 一样简单。
Part Three：怎么做到的——PrivateUse1 与混合 eager/compiled 架构
技术实现上，TorchTPU 踩在 PyTorch 早就预留好的一个扩展点上：PrivateUse1。
PyTorch 的分发器（dispatcher）在 CPU、CUDA 等内建设备之外，留了一个叫 PrivateUse1 的"第三方硬件专用位”。TorchTPU 在分发器中把自己注册为一个新的 C++ 后端，于是开发者拿到手的就是普通的 PyTorch 张量（Tensor）——没有子类（no subclass），也没有包装器（no wrapper）。这一点至关重要：意味着你现有的、写满了原生张量操作的代码几乎不用改，TPU 对上层是"透明"的。
# TorchTPU 之后:普通张量,无包装器,算子调用照常
import torch
import torch_tpu # 导入即完成后端注册(注册到 PrivateUse1)
x = torch.randn(2048, 2048, device="tpu")
w = torch.randn(2048, 2048, device="tpu")
y = x @ w # 走 PyTorch 分发器 -> TPU 后端,与 cuda 上写法一致
y.sum().backward() # autograd 正常,因为它就是普通 Tensor
print(y.device) # tpu:0
但仅注册一个后端还不够，真正的难点是前面说的 eager/编译矛盾。报告的解法是混合即时/编译架构（Hybrid Eager/Compiled Architecture），核心思想是：不追求所有东西都编译，也不退回纯解释执行，而是按算子特性分流——
- 对规整、可融合、形状稳定的计算（典型如大矩阵乘、attention 里的 GEMM），走编译/融合路径，吃满 TPU 的吞吐；
- 对零碎、动态、需要立即返回结果的算子（控制流、动态形状、调试相关），走即时执行路径，保证响应和可调试性；
- 运行时在两者之间做调度，让"eager 扎实"和"编译高效"并存，而不是二选一。
这与推理引擎里"哪些算子融合、哪些算子不融合"的取舍是同一类系统工程：目标从来不是理论上的最优，而是在真实、动态的 workload 下拿到稳定的平均性能。
下面用一段极简 Python 示意这种"按算子分流"的调度直觉（非官方实现，仅帮助理解）：
class HybridRuntime:
"""混合 eager/compiled 的调度直觉:规整重算子编译融合,动态轻算子即时执行。"""
def __init__(self):
self.compiled_cache = {}
def run(self, op, *args):
if self.should_compile(op, args): # GEMM/大注意力:形状稳、可融合
key = (op, tuple(a.shape for a in args if hasattr(a, "shape")))
kernel = self.compiled_cache.get(key)
if kernel is None: # 同形状只编译一次,扛动态形状重编译
kernel = self.compile(op, args)
self.compiled_cache[key] = kernel
return kernel(*args)
return self.eager(op, *args) # 动态/控制流:即时返回,好调试
def should_compile(self, op, args):
big = all(getattr(a, "shape", [0])[-1] >= 512 for a in args if hasattr(a, "shape"))
return op in ("matmul", "attention") and big
def compile(self, op, args): return lambda *a: f"[compiled {op}]"
def eager(self, op, *args): return f"[eager {op}]"
可以看到，缓存的键里带上了形状——同一形状只编译一次，正是为了缓解动态 workload 下"反复重编译"这个老毛病。
Part Four：4917 个模型、98% 跑通——生态兼容性才是真考卷
一个后端"能跑 demo"和"能接住整个生态"完全是两回事。TorchTPU 报告里最有分量的，是覆盖面这组数据：
- 谷歌持续在 TPU 上测试 Hugging Face 上下载量最高的约 5000 个模型；
- 截至报告撰写，覆盖 4917 个、横跨 500 多个架构族；
- 其中 98% 能够跑通不崩溃；
- 在 657 个代表性架构里，97% 的输出与 GPU 结果数值一致。
这两个数字要分开看。“98% 跑通"衡量的是工程兼容性——不报错、不崩、算子覆盖够全；“97% 数值一致"衡量的是数值正确性——TPU 算出来的结果和 GPU 对得上。后者尤其关键，因为不同硬件的浮点实现、归约顺序、低精度格式都可能带来细微差异，敢在 657 个架构上做逐结果比对，说明后端在数值语义对齐上下了硬功夫。
HF Top ~5000 模型
├── 已测 4917 个 / 500+ 架构族
│ ├── 98% 跑通不崩溃 ← 工程兼容性(算子覆盖/不报错)
│ └── 657 代表架构中
│ └── 97% 与 GPU 数值一致 ← 数值正确性(浮点/归约/低精度对齐)
└── 同步接入:vLLM · SGLang · Ray · TorchTitan · Transformers v5
生态侧的"首发名单"也很说明问题：vLLM 和 SGLang 是当下大模型推理的两大事实标准，TorchTitan 是 PyTorch 官方的大规模训练参考实现，Ray 管分布式编排，Transformers v5 则是几乎所有人加载模型的入口。训练、推理、分布式、模型入口四条线同时打通，开发者才有可能真的把生产 workload 搬上去，而不只是在教程里跑个 hello world。
落到工程实践，“设备名可配置"这件事值得认真对待。很多历史代码库把 cuda 写死在各处，迁移时要全局搜索替换、还容易漏。更稳的做法是把设备选择收口到一个入口，让同一份训练/推理逻辑在 CPU、GPU、TPU 之间只靠环境变量或配置切换。下面用 Go 写一个部署侧的设备选择与健康检查示意（真实场景中它常被放在调度器或 sidecar 里）：
package main
import (
"fmt"
"os"
)
// Backend 抽象一个可用的计算后端:设备名只是可切换实现的一部分
type Backend interface {
Name() string
Healthy() bool // 是否就绪/可调度
MemFreeGB() int // 剩余显存/片上内存,粗排资源
}
// 统一的设备入口,绝不把 "cuda" 写死在业务代码里
func pickDevice(prefer string, backends []Backend) (Backend, error) {
for _, b := range backends { // 优先尊重显式偏好,但必须健康
if b.Name() == prefer && b.Healthy() {
return b, nil
}
}
for _, b := range backends { // 偏好不可用时按可用性与资源回退
if b.Healthy() {
return b, nil
}
}
return nil, fmt.Errorf("no healthy accelerator backend available")
}
func main() {
prefer := os.Getenv("TORCH_DEVICE") // 形如 cuda / tpu / cpu,部署期注入
if prefer == "" {
prefer = "cuda"
}
// 同一套逻辑下,接入 TorchTPU 只是多注册一个名为 tpu 的实现
backends := []Backend{
&device{name: "tpu", free: 96, ok: true},
&device{name: "cuda", free: 40, ok: true},
&device{name: "cpu", free: 1 << 20, ok: true},
}
dev, err := pickDevice(prefer, backends)
if err != nil {
panic(err)
}
fmt.Printf("selected device = %s, free = %dGB\n", dev.Name(), dev.MemFreeGB())
}
type device struct {
name string
free int
ok bool
}
func (d *device) Name() string { return d.name }
func (d *device) Healthy() bool { return d.ok && d.free > 0 }
func (d *device) MemFreeGB() int { return d.free }
这段代码想说明的不是某个具体框架的用法，而是一种面向后端中立的架构习惯：业务逻辑面向抽象的"可用加速器"编程，具体跑在哪种芯片上由部署期决定。当 TorchTPU、TorchNeuron 把"新增一个后端"的软件成本压到足够低，这种写法就从"为了优雅而优雅"变成了能直接省钱的能力——你可以在供给宽松、价格更低的硬件上跑同一批 workload，而不必动模型代码。
Part Five：AWS 同步出牌——这不是谷歌一家的单点突围
更值得注意的背景是：TorchTPU 并非孤立事件。在同一个 PyTorch Conference 周期里，亚马逊也同步推进了面向自研 Trainium 芯片的 TorchNeuron 开源后端，同样用 PyTorch 标准设备接口、同样 Apache 2.0，并计划持续向 PyTorch 上游贡献代码。
把两件事并排看，信号就非常清楚了：
PyTorch 原生设备后端(去硬件绑定)
┌──────────────┬──────────────┬──────────────┐
GPU/CUDA TPU Trainium …更多自研芯片
(nvidia 原生) TorchTPU TorchNeuron
(谷歌+Meta) (AWS)
谷歌和亚马逊这两家本身都有自研推理/训练芯片、也都对外卖云算力，它们罕见地选择了同一条打法：用开源软件把"从英伟达 GPU 迁出来"的门槛打下来。它们的共同诉求不是立刻取代 CUDA——那不现实——而是让"换硬件"这件事的边际成本足够低，低到云厂商可以靠性价比和供给把对价格敏感、又被单一供应商绑住的 workload 吸引过来。Meta 的角色也微妙：它既和谷歌合开 TorchTPU，本身又是超大算力买家和 PyTorch 的拥有者，推动框架去硬件绑定，符合它"不被任何一家芯片完全卡死"的长期利益。
这是一场典型的"生态合围”：硬件挑战者们不再各自为战地推自己封闭的软件栈，而是一起把中立框架 PyTorch 做成各方都能原生接入的公地。
Part Six：冷静看待——接口原生 ≠ 无缝迁移，CUDA 仍有厚墙
按惯例把冷水也泼足。
第一，接口兼容不等于性能对齐。 .to("tpu") 能跑，和在 TPU 上跑出有竞争力的吞吐/利用率，是两件事。算子适配的完整性、分布式配置（切分、并行、通信原语）、以及针对性的性能调优，仍要投入真实工程资源。报告自己也强调的是"原生执行能力 + 高覆盖兼容”，而非"全面超越 CUDA 性能”。
第二，CUDA 的护城河远不止一个设备接口。 它包括十几年积累的海量算子库、cuBLAS/cuDNN/NCCL 等高度优化的内核、成熟的 profiler/debugger 工具链、海量踩过坑的工程师、以及大规模训练上被反复验证的稳定性。TorchTPU 撬动的是"入门门槛"和"心理壁垒”——让迁移从"要不要重写一套"变成"先试一把”——但短期内更可能促成的是算力平台多元化，而不是直接替换英伟达。
还可以从"开发者心智"这个更软、但更持久的维度理解这件事。生态的锁定很大程度上发生在习惯层面：当一个工程师的肌肉记忆是 import torch、.cuda()、nvidia-smi，他在选型时会本能地把其他硬件视为"额外成本”。TorchTPU 这类原生后端真正改变的，正是这种本能——当 TPU 也只需要 .to("tpu")、当报错信息和调试器与熟悉的 PyTorch 一致、当 transformers 里一个 device_map 就能切换，“试一下别的芯片"的心理摩擦就降到了几乎为零。心智一旦中立化，采购决策就会更纯粹地按性价比和供给来做，这恰恰是挑战者最希望看到的长期局面。当然，心智迁移需要时间，也需要踩坑经验和社区文档的持续沉淀，但方向一旦确立，它往往比单点性能反超更难逆转——因为它改变的是整整一代新人工程师默认的工作方式。这也是为什么谷歌、亚马逊都愿意在开源软件这种"不直接产生硬件收入"的地方投入重金：它们买的不是当下的性能跑分，而是未来开发者做选择时的那一瞬间的"无所谓”。
第三，TPU/Trainium 主要在各自云上可得。 与可以买卡自建的 GPU 不同，谷歌 TPU 主要通过 GCP 使用，Trainium 绑定 AWS。这对习惯多云或私有化部署的团队，仍是采购和架构层面的现实约束。
用 Python 粗估一下"迁移友好度"如何改变平台选择，能更直观看到 TorchTPU 的战略作用：
def migration_cost(has_native_backend, needs_rewrite, perf_gap, lock_in_concern):
# 原生后端把"重写软件栈"的成本基本清零,只剩调优与采购成本
rewrite = 0.05 if has_native_backend else needs_rewrite
total = rewrite + perf_gap + lock_in_concern
return round(total, 3)
old = migration_cost(False, needs_rewrite=0.6, perf_gap=0.2, lock_in_concern=0.1)
new = migration_cost(True, needs_rewrite=0.6, perf_gap=0.2, lock_in_concern=0.1)
print(f"XLA/封闭栈时代迁移阻力: {old}")
print(f"原生后端时代迁移阻力: {new} # 软件改写这块硬成本被削平")
# 性能差距与采购/绑定仍在,所以是"多元化"而非"替代"
软件改写这块曾经最大的一次性成本被削平后，剩下的差距就是可被时间和投入逐步弥合的工程问题——这正是挑战者最希望看到的局面。
Part Seven：护城河的性质，正在从"锁定"转向"性价比与供给"
把这条新闻拉长看，它是 AI 基础设施权力结构变化的一个标志性切片。
过去十年，英伟达靠"硬件领先一代 + CUDA 生态锁死"建立双重壁垒，软件锁定甚至比硬件性能更难撼动——因为换掉一块芯片是采购决策，换掉一整套软件栈则是组织能力的重写。TorchTPU、TorchNeuron 这类原生后端做的事，本质是把"软件栈重写"从迁移成本里抽掉，让竞争重新回到硬件性能、单位算力价格、供货能力和云服务这些英伟达未必处处占优的维度上。
短期内，英伟达的地位仍然稳固：CUDA 的工具链、人才和大规模验证优势不会因为一个后端开源就消失，AI 训练对顶级互联和成熟生态的依赖也还在。但趋势的方向已经很明确——当 PyTorch 这样的中立框架把 CPU/GPU/TPU/Trainium 拉平成 .to("xxx") 一行之差，“换芯片"就从一个需要立项的工程，逐渐变成一个可以日常比价的选项。
还有一个容易被发布稿忽略、但对实际选型很关键的点：训练与推理的迁移难度并不对称。推理侧的接人通常更快见效——vLLM、SGLang 这类引擎本身就把硬件细节封装在后端里，模型权重和调度逻辑大多可复用，团队最先验证的往往是"同一份权重在 TPU 上能不能以更低的单 token 成本跑起来”。训练侧则重得多：大规模分布式训练涉及数据并行、张量并行、流水并行与集合通信，不同芯片的互联拓扑和通信原语差异显著，TorchTitan 这类参考实现要解决的正是把这套复杂配置在新硬件上重新对齐的问题。因此更现实的迁移路径通常是"先推理、后训练，先离线批处理、后在线长任务"，让团队在风险最小、收益最直观的环节先建立对新后端的信任，再逐步把更核心的训练 workload 挪过去。理解这条节奏，比单纯记住"98% 跑通"这个数字更能帮助决策者排定自己的验证计划。
对开发者和企业，真正务实的姿态不是急于站队，而是从现在起就避免把代码写死在单一后端上：用标准 PyTorch 接口、把设备名做成可配置项、让训练和推理脚本在多种硬件上可切换。TorchTPU 最大的意义，或许正是让"可移植"这件事，从口号变成了一个真的能 pip install 后跑通 4917 个模型的现实。
参考资料：36Kr/量子位等对 PyTorch Conference 2026（圣何塞）谷歌开源 TorchTPU 的报道；技术报告《TorchTPU – A Hybrid Eager/Compiled Architecture for Native PyTorch Execution on TPU》arXiv:2610.24161；项目地址 github.com/google-pytorch/torch_tpu；DOIT 关于亚马逊 TorchNeuron 同步开源的快讯。文中 4917 模型/98% 跑通/657 架构 97% 数值一致、Anthropic 百万 TPU、Meta 数十亿美元租用等数据来自公开报道与技术报告；代码为原理示意，非谷歌官方实现。