AIREITER把 PyTorch 代码从 cuda 迁移到 tpu，表面上改动并不大；但 TPU 是否可用、依赖包是否匹配、编译过程以及张量形状，最终都会决定这次迁移能否顺利、性能是否理想。
TorchTPU 到底解决了什么问题
TorchTPU 是面向 Google TPU 的原生 PyTorch 后端，不是托管模型 API，也不是 TPU VM 的替代品。PyTorch 官方公告提到，它支持 eager execution、torch.compile、分布式原语以及大约 450 个 ATen 算子；公开代码仓库则列出了 ATen 内核、编译缓存、分布式示例和 torch.compile 集成。
迁移时，代码层面的目标非常直观：
# Before
model.to("cuda")
# First TPU port
model.to("tpu")不过，这只能说明代码具备可移植性，不能保证性能。PyTorch 的公告也明确指出，TPU 能否跑出峰值性能，仍取决于模型形状、分片方式以及是否使用适合该硬件的算子。
安装前需要准备什么
只有在已经配置好 Google TPU 的环境中，TorchTPU 才真正有用。代码仓库建议在 TPU VM 上运行其 notebook，并将 v6e Google Compute Engine 实例列为最容易上手的环境。只配备 CPU 的笔记本可以安装相关软件包，但无法验证 TPU 执行。
公开 README 支持 Python 3.11+，并指定了 torch_tpu==0.2.0。建议使用隔离环境，避免实验意外继承 CUDA 版本或其他无关依赖。
安装 torch_tpu，避免带入 CUDA 依赖
先创建虚拟环境，再通过指定的 TorchTPU 版本和 CPU 版 PyTorch wheel 完成安装：
python -m venv .venv
source .venv/bin/activate
pip install torch_tpu==0.2.0 \
--extra-index-url https://download.pytorch.org/whl/cpu按照 README 的说明，如果省略这个额外索引，pip 可能会拉取 CUDA 版 PyTorch，这并不适合 TPU VM。
第一次迁移建议直接使用已发布的 wheel。源码构建则需要 Bazel 或 Bazelisk，并按照仓库文档使用 --config=wheel_common；必要时还要加上 --config=no_rbe。
用最小 PyTorch 程序验证后端
请在 TPU VM 上运行下面的代码：
import torch
x = torch.randn(2, 2, device="tpu")
y = torch.randn(2, 2, device="tpu")
z = x @ y
print(z)
print(z.device)如果程序成功输出结果，且设备显示为 tpu，就说明张量分配和矩阵乘法已经进入 TPU 后端。如果设备不可用，先检查 VM 运行时和软件包安装情况，不要急着修改模型代码。
TorchTPU 同时支持便于调试的 eager execution，以及用于编译执行的 torch.compile。PyTorch 公告称，在项目测试中，DeferAndFuse 模式相比严格 eager 模式的速度提升可以达到 50% 或更高；这属于项目测试结果，不应视为适用于所有模型的通用基准。
把真实的 PyTorch 工作负载迁移过来
最稳妥的做法是先完成最小改动，确认单设备路径没问题后，再加入 TPU 专用的分布式配置。
- 将设备选择改为
torch.device("tpu")；如果环境支持，也可以使用torch.accelerator.current_accelerator()。 - 确保模型和输入张量位于同一个设备上。
- 在启用编译或分布式执行前，先跑一个短 batch。
- 测量性能时保持 batch size 和张量形状稳定；形状变化可能触发额外的编译工作。
- 在投入更大规模迁移前，先测试具有代表性的模型路径和分布式模式。
使用 CUDA 的分布式代码通常依赖 NCCL，而 TorchTPU 发布的示例使用的是 TPU 专用进程组后端：
import torch
import torch.distributed as dist
device = torch.device("tpu")
dist.init_process_group(backend="tpu_dist")
model = model.to(device)具体的分布式启动命令取决于 TPU 环境和工作负载。仅仅把设备字符串改掉，并不会自动完成分片、进程放置或输入流水线配置。
第一次运行的表现很可能与后续运行不同。TorchTPU 提供了文档化的编译缓存，可以复用重复出现的计算图，因此应将预热时间与稳定状态下的吞吐量分开记录。
TorchTPU 已经能做什么，还有哪些地方不够成熟
Google 已公布 TorchTPU 与 vLLM、SGLang、TorchTitan 以及 Hugging Face 工具链的集成。其推理公告介绍了 vLLM 和 SGLang 的原生 TPU 路径，包括 KV cache 搬运和服务能力。Google 的Cloud 公告则报告了 TorchTitan 运行 Qwen 3 8B 时 52% MFU 的结果，以及相关生产应用情况。
但从实际落地角度看，TorchTPU 的成熟度仍没有公告标题看起来那么全面：
| 领域 | 现有资料能够确认的内容 | 你仍需要自行验证的事项 |
|---|---|---|
| 安装 | Python 3.11+、torch_tpu==0.2.0、CPU PyTorch 索引 | 当前 VM 镜像能否与 wheel 顺利解析并完成安装 |
| 可移植性 | eager mode、torch.compile、tpu 设备以及分布式支持 | 不支持的算子，以及针对具体模型的回退路径 |
| 性能 | 大约 450 个 ATen 算子，以及 TorchTitan 运行 Qwen 3 8B 时项目报告的 52% MFU 结果 | 你的模型在预热时间、形状、batch size 和成本方面的表现 |
| 开发流程 | 公开示例、notebook、调试工具和基准测试工具 | TPU 访问条件、Bazel 的复杂度以及可用的可观测性工具 |
| 当前缺口 | 功能覆盖仍在持续完善 | 你的训练栈是否依赖尚未针对该工作负载公开文档的功能 |
是否采用 TorchTPU，可以按这条标准判断
如果你已经拥有 Google TPU 访问权限，运行的是 PyTorch 模型，并且希望在训练或推理中使用统一的框架路径，那么现在就值得测试 TorchTPU。它更适合愿意分析性能、调整张量形状的团队，而不是期待只改一行代码就获得与 GPU 相当效率的用户。
如果你的关键工作负载需要成熟的兼容性矩阵、可预测的冷启动表现，或者依赖 TorchTPU 尚未提供的某项功能，那么暂时不要急着迁移。可以先选定一个模型、一个 TPU slice 和一套固定评估数据，分阶段记录安装耗时、首次编译时间、稳定状态吞吐量以及输出正确性。
TorchTPU 常见问题
TorchTPU 和 PyTorch/XLA 是一回事吗？
不是。TorchTPU 是较新的原生 PyTorch 后端，围绕 eager execution 和 torch.compile 设计；PyTorch/XLA 则是更早的 TPU 集成路径。现有 XLA 代码需要单独评估迁移工作量。
没有 TPU 的笔记本可以安装 TorchTPU 吗？
可以安装软件包，但没有兼容的 Google TPU 基础设施，就无法验证 TPU 执行。官方代码仓库给出的实际配置前提是 TPU VM。
为什么 TorchTPU 第一步运行得很慢？
第一次执行可能需要编译计算图并填充缓存。排查性能时，应单独测量后续迭代，并在诊断重新编译问题时保持张量形状稳定。
把 cuda 改成 tpu，就能保证性能一样吗？
不能。替换设备只是完成了正确性迁移；在比较性能前，还需要分析张量形状、分片方式、算子覆盖情况和预热过程。
TorchTPU 已经适合生产环境了吗？
Google 已介绍 TorchTPU 在生产推理和训练中的应用，但公开生态组件和具体功能的成熟度并不完全相同。在正式采用前，应针对实际模型、服务引擎和分布式模式进行测试。
在 TPU VM 上安装指定版本的 wheel，完成 2×2 验证，然后用一个已经预热的工作负载与当前 GPU 路径进行对比。这样可以分别看清代码可移植性、编译器预热、算子覆盖范围以及硬件成本。