TorchTPU 迁移实战：3 个 CUDA 卡点

我把手上一个典型的 GPU 训练工程搬到 TorchTPU 上试了一遍，「换芯片要重写几个月代码」这句话，在 2026 年 10 月 20 日之后需要打个折。那天 PyTorch Conference 上，Google 与 Meta 的 PyTorch 团队一起把 TorchTPU 开源了：TPU 成为一个原生 PyTorch 设备，理论上你只需要把 .to("cuda") 改成 .to("tpu")。我把手上一个典型的 GPU 训练工程拿去试了一遍，改一行是真的，跑起来之后才发现问题都在那一行之外。这篇把我实际撞到的 3 个卡点和一份能直接跑的审计脚本放出来。

一、TorchTPU 到底是什么

先把范围说清楚，避免误判（来源：Google Cloud 官方博客与 GitHub 仓库 README）：

2026 年 10 月 20 日 PyTorch Conference 2026（圣何塞）宣布开源，由 Google 与 Meta 的 PyTorch 团队共同开发，协议 Apache 2.0；
版本 v0.2.0 已上 PyPI，要求 Python 3.11 及以上；
通过 PyTorch 的 PrivateUse1 扩展接口接入，张量是普通 PyTorch 张量，没有子类、没有包装器，保留视图与原地修改语义；
官方定位是取代 PyTorch/XLA，新负载推荐直接用它；
配套仓库四个：TorchTPU（核心后端）、vLLM-TorchTPU（与 Inferact 合作）、SGLang-TorchTPU（与 RadixArk 合作）、TorchTitan（训练配方）。

执行模式有三种 eager：Strict（默认，异步分派，标准 PyTorch 语义）、Debug（同步执行并做张量越界检查）、DeferAndFuse（融合连续的 eager 算子以减少 HBM 流量）。要跑满吞吐则走 torch.compile(backend="tpu")，由 TorchDynamo 捕获图，经 StableHLO 降到 OpenXLA 生成 TPU 机器码，编译产物可跨进程、跨主机、跨任务重启缓存。

二、卡点 1：设备写死的地方比你想的多

.to("cuda") 只是冰山一角。真实工程里设备假设散落在各个角落，我把它们整理成了可机器判定的规则。

python 复制代码

# 典型写法，一眼看不出问题
model = MyModel().cuda()
x = torch.randn(2, 3).to("cuda:0")
torch.cuda.synchronize()

改成设备无关之后是这样：

python 复制代码

device = torch.device(cfg.device)      # "cuda" 或 "tpu"，从配置读
model = MyModel().to(device)
x = torch.randn(2, 3).to(device)
# 同步语义交给框架，不要显式调用 torch.cuda.*

torch.cuda.manual_seed 这类调用容易被漏掉，它在没有 CUDA 的环境里会直接抛错，属于硬阻断。判断标准很简单：凡是出现 cuda 字面量、.cuda() 或 torch.cuda.* 的地方，都算一次迁移成本。

三、卡点 2：分布式那一行比设备那一行更要命

这一条是我在迁移里花时间最多的。GPU 上写的是：

python 复制代码

torch.distributed.init_process_group(backend="nccl")
world = torch.cuda.device_count()

TorchTPU 用的是 tpu_dist 进程组（来源：Google Cloud TPU 官方文档），支持 DDP、FSDP2 与 DTensor，并有面向 TPU 互联拓扑的 mesh。但真正的问题不在 backend 名字，而在进程与设备的对应关系 ：TPU 侧是一进程一设备，而且官方明确写了 TPU 设备当前只支持一个用户态进程。也就是说，用 GPU 数量去推断 world_size 的那套写法会直接错。

还有一处隐性成本（来源：PyTorch 官方博客对 TorchTPU 的限制说明）：官方与 PyTorch 博客都提到，Skinny Collectives（分布式里的细粒度集合通信）相比融合模块有 30% 到 40% 的性能损失，原因是 XLA 无法重叠计算与通信。换句话说，通信写得越碎，在 TPU 上罚得越重，而在 GPU 上你可能根本没感觉到。

四、卡点 3：自定义 CUDA 算子是硬阻断

这一条没有变通余地。通过 torch.utils.cpp_extension、load_inline 或 CUDAExtension 加载的自定义算子，在 TPU 上没有对应实现。官方给的替代路径是 Pallas 自定义核、预置的 Tokamax 核与 Helion tile 核，它们可以作为标准 PyTorch 函数被调用。

如果你的工程重度依赖某个手写 CUDA 算子，那么迁移的第一步不是改设备，而是先把这个算子用 Pallas 重写一遍。这一步的工作量，跟框架本身无关。

五、把这三件事做成可复现的检查

我写了个静态扫描脚本，扫的是「代码里有多少 CUDA 假设」，不是「TPU 上能跑多快」。判据三态：BLOCKER（必须改）、REVIEW（要人看一眼）、OK。

bash 复制代码

python tpu_migration_audit.py --selftest
python tpu_migration_audit.py --dir /path/to/your/project

在一个典型 GPU 训练工程上的实际输出：

复制代码

扫描 1 个文件：BLOCKER 1 / REVIEW 0 ⇒ 先清 BLOCKER
  [BLOCKER] /tmp/tpu_demo/train_demo.py
      L   2 R4_custom_cuda_ext 加载自定义 CUDA 算子 → 需要 Pallas / Tokamax / Helion 重写
      L   4 R3_nccl_backend 分布式进程组写死 nccl → TorchTPU 用 tpu_dist 进程组
      L   5 R8_gpu_count_assume 用 GPU 数量推断 world_size → 一进程一设备，必须改
      L   6 R2_cuda_call 直接调用 .cuda() 或 torch.cuda.*
      L   7 R5_cudnn_benchmark GPU 专属的后端开关 → TPU 上无对应语义

设计上有几条纪律值得说明。其一，注释里的示例代码不计入 ，否则满屏假红，判据会被训练成忽略。其二，输入为 None 时判 REVIEW 而不是 OK ，因为说不清就不该放行，这与 KV cache 检查用的是同一条规矩。其三，只有 BLOCKER 阻断，REVIEW 不阻断，比如 pin_memory=True 和 non_blocking=True 在 TPU 上是无效但无害的优化，判 REVIEW 就够了，把它判成阻断会让人直接放弃这套检查。

需要交叉验证的是：静态扫描只能看到 .py 文本。算子写在 .cu 文件里它看不到，动态拼出来的设备字符串它也看不到，所以扫干净不等于没问题，这只是一道下限检查。

六、性能数字该怎么读

先列数字，再说怎么读（以下均为 Google 及合作方自报，无独立第三方复现）：

TorchTitan 在 Ironwood 上训练 Qwen 3 8B 达到 52% MFU（来源：Google 官方博客，厂商自报），128 到 1,024 芯片区间保持稳定扩展；
Llama 3 405B 在 128 到 1,024 片 Ironwood 上约 57% MFU；
一个微调案例里，单步时间从开箱即用的 11.49 秒降到调优后的 2.78 秒，其中 DeferAndFuse 单独贡献 1.15 倍，torch.compile 完整路径贡献 3.80 倍；
约 5,000 个 Hugging Face 模型在 eager 与编译两种模式下持续测试，Tier-1 无崩溃通过率 98%，与 GPU 执行的数值一致性 97%；
SemiAnalysis InferenceX 的基准被引述为：TPU v7 配 vLLM-TorchTPU 在特定推理任务上达到每百万 token 2.01 美分。

读法上的两条提醒。第一条，开箱即用与调优后差了 4 倍多 （11.49 秒对 2.78 秒），这意味着「改一行配置」换来的是「能跑」，不是「跑得好」，两者之间隔着一整套硬件感知的调优手册。第二条，官方与 PyTorch 博客都明确写了：模型 shape、切分方式和算子往往是按 GPU 架构调优过的，换个设备能跑，但跑到峰值需要重做这些假设；再加上 TPU 指令集架构并不公开，这条路并不像接口层看起来那么平。

七、辩证：这不等于 CUDA 的工具链被追平

另一个角度必须给出。开放的是接口层，不是配套体系的对等。CUDA 的护城河来自十余年的算子库、profiler、踩坑沉淀和社区答案；TorchTPU 解决的是「迁移成本」这一项，不是「工具链厚度」这一项。

也不必把它读成谁要取代谁的现实判断。同一周里，Anthropic 向 Nscale、Lambda 租用英伟达芯片的合约上限被披露到约 860 亿美元（合同延续到 2033 与 2034 年），亚马逊在前一天发布了自家的 TorchNeuron 迁移工具。云厂商前后脚做同一件事，说明的是迁移成本正在被当成竞争变量，而不是某一家已经出局。

对开发者来说，落到执行层面只有一条：**保持 PyTorch 标准接口的干净度，就是未来的议价权。**少写设备字面量，少绑死 backend，少依赖手写算子。今天省下的这几个小时，换的是明天能换硬件的自由。

数据与事件来源

Google Cloud 官方博客《TorchTPU is now open source》，2026-10-20
GitHub 仓库 TorchTPU，Apache 2.0，v0.2.0 on PyPI，要求 Python 3.11 及以上
Google Cloud TPU 官方文档（三种 eager 模式、torch.compile 路径、tpu_dist 进程组）
MarkTechPost、AppStack Insider、Forkast 等媒体报道（性能数字均注明为厂商自报）
SemiAnalysis InferenceX 基准（经 PyTorch 团队引述）
本文所有可复现检查见同目录 tpu_migration_audit.py；--selftest 共 10 项（正控 4、负控 6）全部通过

你怎么看？如果你的工程现在就跑一遍这个扫描，报出来的 BLOCKER 会有几条？欢迎在评论区把结果贴出来，我想看看真实的分布。

互动思考题：

你的代码里 device 是从配置读的，还是散落在各个文件里的字面量？
如果换一次算力的成本从几个月降到几天，你会因此改变框架的选型策略吗？
静态扫描扫不出的动态设备字符串，你会用什么办法兜住？

欢迎在评论区留下你的扫描结果，我会挑几条跟着复现。