PyTorch 原生跑上 TPU 了

今天刷到 Hugging Face 团队发的一篇博客，说 TorchTPU 这个原生 PyTorch 后端正式落地了。简单讲，以前想在 TPU 上跑 PyTorch，要么把代码迁到 torch_xla，要么干脆去学 JAX，门槛劝退了不少人；现在 TPU 在 PyTorch 里就是另一个 device，把 Hub 上的模型直接 .to("tpu") 就能跑。

这次不只是「能跑」，而是生态层面的原生支持：Transformers、Diffusers、Kernels 都直接适配了 TPU。也就是说文本模型、图像和视频生成模型都能在 TPU 上原生推理，不用再为某个算子单独写 JAX 版本。对做多模态应用的团队来说，这个吸引力比单纯的 benchmark 数字大得多。我的判断是，这件事真正的意义在于「降低迁移成本」。TPU 的算力性价比在训练和批量推理上一直有竞争力，但过去卡在软件栈上，很多团队评估完迁移工作量就放弃了。现在 PyTorch 原生支持，等于把评估周期从几周压到几天，Google Cloud 上的 TPU 会因此多出一批原本不会考虑它的用户。当然也别急着 all in。新后端刚发布，算子覆盖、分布式训练成熟度、调试工具链都还需要时间验证，博客里也提到团队自己做了大量测试来保证可用性。我的建议是：新项目可以拿小模型先在 TPU 上跑一轮推理压测，老项目继续观望一两个版本再决定是否迁移。对跨境和出海团队来说，这还多了一层意义：如果业务主要跑在 Google Cloud 上，TPU 原生支持意味着推理成本有机会再压一档，尤其是图像、视频类生成场景。接下来我会关注两件事，一是社区反馈的踩坑记录，二是主流微调框架什么时候跟进适配。