Hugging Face ❤️ TorchTPU：原生支持TPU

Hugging Face 于 10 月 20 日宣布，PyTorch 代码现在可以直接跑在谷歌 TPU 上了。这项名为 TorchTPU 的新后端把 TPU 变成了 PyTorch 中一个普通设备，开发者只需在加载模型时指定 device_map=“tpu”，就能把 Hub 上的模型直接部署到 TPU，不再需要把代码移植到 torch_xla，也不用转去学 JAX。Transformers、Diffusers 和 Kernels 三大库均已原生支持。 这件事的背景是 TPU 长期存在的软件短板。TPU（张量处理单元）是谷歌自研的 AI 加速芯片，设计至今约十年，Gemini 等自家模型都靠它训练和推理，谷歌云上从单芯片到数千芯片的集群均可租用，算力性价比在训练和大规模推理场景中颇具吸引力。但 TPU 的软件栈一直围绕 TensorFlow 和 JAX 生长，PyTorch 用户想用就得走 torch_xla 这条老路：张量是惰性求值的，代码不按预期时机执行，小改动可能触发计算图重编译，调试起来相当费劲。Hugging Face 此前推出的 Optimum TPU 虽然补上了模型支持，但终究是绕路方案。TorchTPU 落地后，Optimum TPU 不再是运行 Transformers 和 Diffusers 模型的必需品。 对跨境电商卖家和出海团队来说，实际影响在于推理成本的选择空间变大了。此前用 PyTorch 生态做商品图生成、多语言文案、客服模型微调的团队，基本被锁在英伟达 GPU 上；现在同一套代码可以平移到谷歌云 TPU，在算力紧张或 GPU 价格高企时多了一个议价筹码。Diffusers 的原生支持尤其值得关注：Stable Diffusion 3.5 这类模型只需 pipe.to(“tpu”) 加一次 compile 调用即可出图，首次调用完成编译后，相同尺寸的图片会复用计算图；模型大到单芯片放不下时，传入 TensorParallelConfig 就能在加载时自动把 transformer 分片到多颗芯片上。安装同样简化为一句话：pip install torch-tpu，并附加 CPU 版 torch 的索引地址，避免 pip 误拉体积庞大的 CUDA 版本。模型默认以 eager 模式运行，追求性能可再启用 torch.compile 或分布式配置。