10 月 20 日的 PyTorch Conference 2026 上，Google 宣布 TorchTPU 正式开源，Hacker News 上 242 个赞的讨论指向同一个仓库：google-pytorch/torch_tpu。路透社此前的独家报道给这次发布定了性：Google 联手 Meta 的 PyTorch 生态，目标是削弱英伟达赖以垄断的软件护城河。硬件军备竞赛打到现在，连 Google 都承认 TPU 普及的真正瓶颈是软件入口，而这次开源就是冲着这个入口来的。本文基于仓库代码、官方文档与开发者博客，把这套栈的架构、上手路径与成熟度讲清楚。
它解决什么问题
TPU 驱动着 Google 自家的 Gemini 与 Veo，也服务云上客户，但长期以来，想把 PyTorch 工作负载迁到 TPU 的团队都要过一道弯：PyTorch/XLA 路线提供的是相对薄的外层集成，写法与调试体验和原生 CUDA 路径有明显落差。TorchTPU 的设计原则只有三条：用起来就应该像 PyTorch，现有训练循环只改初始化就能跑；以最小代码改动迁移既有负载；提供榨取 TPU 全部算力的 API 与工具。
实现路径基于 PyTorch 的 PrivateUse1 接口做深度集成，没有子类、没有包装器，TPU 张量就是普通 PyTorch 张量，tensor.to('tpu') 的体验对齐 tensor.to('cuda')。需要说明的是，仓库实际创建于今年 3 月，本次发布的是公开开源形态，许可证为 Apache 2.0，PyPI 包 torch-tpu 已发布 0.2.0 版本。
架构：Eager 优先，编译直通 StableHLO
TorchTPU 最值得研究的是它的执行模型，官方称之为 Eager First，提供三种 Eager 模式。Debug Eager 逐操作调度、每步与 CPU 同步，慢但适合排查形状不匹配、NaN 与显存溢出。Strict Eager 单操作调度但异步执行，贴近默认 PyTorch 体验。Fused Eager 对操作流做自动反射分析，动态融合成更大计算块以提升 TensorCore 利用率、降低显存带宽开销，官方实测比 Strict Eager 快 50% 到 100% 以上，且无需任何配置。三种模式共享一套编译缓存，可单机或跨多主机持久化，省掉重复编译。
静态编译路径同样有明确取舍：torch.compile 下 Dynamo 捕获 FX 图后，不走 Torch Inductor，而是直接以 XLA 为主后端，翻译成 StableHLO 再生成 TPU 二进制。官方给出的理由很直接：XLA 针对 TPU 拓扑久经考验，原生理解密集计算与 ICI 集合通信的重叠优化。TPU 系统的硬件形态也体现在文档里：芯片通过 ICI 互联组成 2D 或 3D torus，片内分 TensorCores 负责密集矩阵、SparseCores 负责嵌入、聚集散布与集合通信。文档甚至给出了一条硬件感知建议：把常见的 64 维注意力头改成 128 或 256，以匹配当代 TPU 的峰值矩阵乘效率。
分布式方面，DDP、FSDPv2 与 DTensor 开箱支持，大量基于 PyTorch 分布式 API 的第三方库无需改动。一个差异化能力是 MPMD，即多程序多数据执行，允许不同 rank 运行不同代码，例如 rank 0 额外做日志与分析，必要时隔离通信原语保证正确性，同时尽可能保留 XLA 的通信与计算全局重叠。自定义内核目前支持 Pallas 与 JAX 写法，通过 @torch_tpu.pallas.custom_jax_kernel 接入，更上层的 Helion DSL 在路线图中。
上手路径
官方安装方式很直接，在 TPU 虚拟机上建好虚拟环境后执行：
pip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpu注意这个额外的索引源提供的是 CPU 版 torch，不加会装到 CUDA 版本，这是文档里专门强调的坑。Python 版本要求 3.11 以上，与 PyTorch 的发布兼容矩阵对齐。文档站在 torchtpu.readthedocs.io，交互式笔记本用 Marimo 编写，覆盖 Llama、Qwen、ResNet、DLRM、minGPT 与 UNet 等模型的分布式训练与推理示例，官方建议用 v6e 规格的 GCE 实例最省事。从源码构建走 Bazel，需要 --config=wheel_common，没有远程构建执行凭据时追加 --config=no_rbe，且后者必须写在后面才会生效。
与 PyTorch/XLA 的正面对照
把两者的差异列成表更清楚。并行模型上，XLA 路线长期只支持纯 SPMD，TorchTPU 增加 MPMD 分歧执行；集成深度上，从薄外层封装换成了 PrivateUse1 深度集成；编译路径上，Dynamo 捕获后跳过 Inductor 直映射 StableHLO；Eager 体验上，从较弱支持升级为三模式加自动融合；自定义内核从有限支持变为原生接入 Pallas 与 JAX；分布式 API 从受限变为开箱支持三件套。可以看出，这不是一次增量优化，而是把 PyTorch 侧的原生体验按 TPU 硬件特性重新做了一遍。
成熟度评估
路线图还剩几块硬骨头：预编译 TPU 内核库以降低首次执行延迟、torch.compile 动态形状的一等公民支持（通过 XLA 的有界动态机制推进）、原生多队列以适配计算与显存流解耦的异步代码库，以及与 vLLM、TorchTitan 的深度集成并验证到完整 Pod 规模的线性扩展。对生产团队，判断标准因此很清晰：训练路径与主流分布式 API 现在就能用；动态形状负载、首次编译延迟敏感的场景，建议等内核库与动态形状落地后再上。对研究团队，Fused Eager 的自动融合与 PrivateUse1 的集成方式，本身就是值得细读的工程样本。
把视角拉到产业层，这次开源与 Cloudflare 收购 Deno、vLLM 生态的扩张是同一条叙事线的不同切面：推理与训练栈的竞争焦点，正从模型本身转向「模型跑在谁的地基上」。Google 把 PyTorch 入口亲手交出来，等于承认了一件事：在英伟达的软件护城河面前，再强的自研芯片也需要借道竞争对手主导的框架生态。对开发者这是纯粹的好事，TPU 的门槛从「会写 XLA」降到「会写 PyTorch」，云上算力的比价维度又厚了一层。