PyTorch 原生跑 TPU 了

今天看到 Hugging Face 团队发的一篇博客，讲他们做了个叫 TorchTPU 的东西，让 PyTorch 能原生跑在 TPU 上。以前想用 TPU，要么把代码迁到 torch_xla，要么干脆去学 JAX，迁移成本劝退了不少人。现在官方说法是 TPU 变成「另一个 device」，模型从 Hub 拉下来直接 .to("tpu") 就行，Transformers、Diffusers、Kernels 都原生支持。

TPU 是谷歌自研的 AI 加速器，专门为神经网络里的矩阵运算设计，做了大概十年，Gemini 就是用它训练和推理的。谷歌云上从单芯片到多芯片集群都能开，这点对中小团队其实挺友好，不用一次性囤卡。我的判断是，这件事的意义不在性能参数，而在生态门槛。PyTorch 是绝大多数算法同学的主场，TPU 过去更像「谷歌内部+JAX 圈子」的玩具。现在 PyTorch 代码几乎零改动就能上 TPU，等于把谷歌的算力摆到了主流开发者面前，多云议价的空间也大了。当然别急着 all in。新后端早期坑肯定不少，算子覆盖、分布式训练、调试工具链都得打个问号，博客里也提到他们自己做了大量测试。我的建议是先拿推理和微调场景试水，训练主链路再观望一两个版本。对做出海 SaaS 的团队来说，这算个值得记一笔的信号：算力选择正在从「英伟达 or 其他」变成「英伟达 + 多家加速器」。早点把代码里的设备抽象做干净，以后换芯片才不至于伤筋动骨。