PyTorch 原生跑上 TPU 了

今天刷到 Hugging Face 那篇 TorchTPU 的文章，第一反应是：终于不用为了用 TPU 去学 JAX 了。以前想把 PyTorch 模型搬上 TPU，要么改代码走 torch_xla，要么整个团队切换心智模型去写 JAX，迁移成本高到很多项目直接放弃。现在 TPU 在 PyTorch 里就是另一个 device，模型从 Hub 拉下来直接 .to("tpu") 就能跑，这个体验变化是实打实的。

文章里提到 Transformers、Diffusers 和 Kernels 都已经原生支持 TPU。这意味着文本模型、图像和视频生成模型都能直接在 TPU 上推理和微调，不用再维护一套 XLA 专用分支。对做多模态产品的团队来说，少一层适配就少一堆线上玄学问题，这点比跑分更让我在意。我的判断是，这件事真正的意义不在技术炫技，而在采购决策。过去选 TPU 往往意味着绑定 JAX 生态，团队得重新招人、重新积累；现在 PyTorch 代码几乎零改动，TPU 就变成了和 GPU 并列的一个选项。云上单卡到多卡都能开，做推理成本对比的时候终于可以把 TPU 放进同一张表里算账了。当然也别急着 all in。原生后端刚出来，算子覆盖、分布式训练成熟度、调试工具链都还需要时间验证，文章自己也说做了大量测试来兜底。我的建议是先在推理侧或小规模微调上试水，把吞吐和成本跑出真实数据，再决定要不要把训练主链路迁过去。总的来说，TPU 从「另一个生态」变成「另一个设备」，对出海做 AI 应用的团队是利好，供应链上多了一个能压价的筹码。接下来我会关注两件事：社区里真实的踩坑反馈，以及各家云厂商对 TPU 实例的定价会不会跟着松动。