代码里把运行设备从 cuda 改成 tpu，表面上只像换个插座。
2026 年 10 月 20 日，Google 在 PyTorch Conference 上把 TorchTPU 推到 0.2.0 版本，并定于 11 月 15 日归档服役多年的旧版 PyTorch/XLA。就在前一天，Amazon 也将适配自研 Trainium 芯片的 TorchNeuron 推入公测。两大云厂商在 48 小时内亮出同一套策略：不再逼开发者学私有框架，而是借 PyTorch 的标准机制，把算力流量从英伟达 GPU 往自家芯片上引。
不用重写业务代码的承诺，给苦于硬件采购预算的企业画了一张大饼。但在现实工程里，把算力调度的指令改掉只是第一步，底层编译器和手写算子的断层远没有公关稿写得那么平整。
标准接口分流与断舍离
过去几年，PyTorch 开发者调用 TPU 体验极差。旧的 PyTorch/XLA 依赖惰性张量模式，必须塞进各类私有同步指令，随时伴随着未知的黑盒重编译。
这次 Google 和 AWS 统一了突围姿态。双方均采用 PyTorch 的 PrivateUse1 扩展机制，把设备名直接收归为原生的 tpu 与 neuron。开发者依赖即时执行（Eager mode）配合 torch.compile，由 TorchDynamo 抓取图结构，再经由 StableHLO 对接各自的后端编译器。
与架构调整同步发生的，是激进的技术债务出清：
- 运行环境锁定在 Python 3.11 及以上，不再兼容老旧环境；
- 维护重心转向 TPU v5e、v5p、Trillium（v6e）以及下一代 Ironwood（v7x）；
- 彻底抛弃早期硬件，TPU v2、v3、v4 移出支持列表；
- 旧版 PyTorch/XLA 彻底归档，不再接受常规特性迭代。
抛掉历史包袱换来了亮眼的基准数字。在 8,192 块 Ironwood TPU 集群上，运行 Qwen 3 8B 拿下了 52% 的 MFU。针对 FLUX.1 的优化链路里，单块 TPU v6e 单步耗时压到 122.51 毫秒；在推理引擎 vLLM 上，Qwen 3.5 的运行成本被标定在每百万 Token 约 2.01 美分。
静态计算图的顺境与动态业务的暗礁
官方测试集里的高通过率，前提是模型结构规整、张量形状固定。
只要代码踏出静态边界，底层的 XLA 编译器依然容易卡壳。Eager 执行模式没有改变 XLA 需要提前推导静态图的事实。面对带有动态切片、变长上下文或复杂控制流的代码，系统往往反复触发重编译：
- 在 Whisper 风格的任务中，一旦张量形状改变导致未命中缓存，初次编译耗时可达 18.93 秒，而后续单步执行只需 0.02 秒，相差近千倍；
- 在 Stable Diffusion 3 去噪流程里，张量 View 存储偏移量稍有变动便诱发重编译，以至于生态库不得不建议用户关闭动态特性；
- 基础数学算子依然存在缺口，缺少
torch.linalg.cholesky_ex与aten::unique_dim，遇到相关逻辑往往退回 CPU 或报错中断； - 编译模式下调用缩放点积注意力时，反向传播偶尔出现浮点溢出，让长时间训练中断。
Hugging Face 测试套件里高达 96.5% 的运行通过率，主要验证了主干网络的前向通路。但在生产训练中，多变的长文本输入和复杂的批处理机制，极易让脆弱的动态 Shape 支持暴露出吞吐断崖。
算力采购单与真实的迁移负债
云厂商正面叫板英伟达的抓手是定价。算力预算不是技术情怀，账面差距确实一目了然。
Trillium 单卡时租 2.70 美元，AWS 宣称 Trainium2 相对 H200 具备 30% 到 40% 的性价比空间。对开销庞大的团队来说，这类数字极具吸引力。
但分布式训练深水区容不下纯粹的财务算术。TorchTPU 0.2.0 目前仍未提供完整的 torch.Stream 抽象，开发者无法精细编排通信与计算的流水线重叠。FSDPv2 和 MPMD 在文档中仍标记为实验状态，生产集群多需依靠定制化方案兜底。
更沉重的成本在于算子积累。如果现有模型架构严重依赖 Triton 或手写 CUDA 核函数，要发挥 ASIC 芯片的标称算力，就必须用 Google 的 Pallas 或 AWS 的 NKI 彻底重写。这不仅需要熟悉硬件架构的高水平系统工程师，还会直接拖慢迭代节奏。
《韩非子》有云：“兵不完利，与操策者异。”兵器不够齐备，车夫驾车失控的风险就陡增。
换掉硬件不等于自然卸下包袱。当前的边界相当明确：如果是标准化、结构固定的大并发推理业务，切到 TPU v6e 这类专用芯片确实能省下预算；但若是高度依赖自定义算子与前沿架构的预训练团队，继续支付英伟达的软硬件确定性溢价，依然是更理性的工程选择。