Hugging Face发布TorchTPU：PyTorch原生跑上TPU

Hugging Face于10月20日发布博客，宣布TorchTPU：一个让PyTorch原生支持TPU的后端。按照官方说法，Transformers、Diffusers与Kernels现在都能在TPU上原生运行，选一个模型、把设备设为“tpu”即可。

在此之前，PyTorch用户想用TPU，通常要移植到torch_xla，或者干脆改学JAX。XLA路线的张量是惰性的，代码不总在预期的时刻执行，细微改动就可能触发图重编译，调试体验也不直观。Hugging Face曾用Optimum TPU弥补这一缺口，提供推理与训练支持。

TorchTPU改变了这个前提。它把TPU作为PyTorch的新后端接入，用户写普通PyTorch代码即可，附带为TPU调优的Pallas内核、分布式训练支持，以及可选的torch.compile集成。在TPU虚拟机上安装只需一条pip命令，并指定CPU版torch的索引地址。

具体用法被压到最小：Transformers v5中传入device_map="tpu"即可把模型直接加载到TPU；Diffusers中pipe.to("tpu")完成推理，需要性能时再对transformer调用compile并指定tpu后端。官方表示，随着原生支持到位，运行Hugging Face模型已不再需要Optimum TPU。

对显存放不下的模型，TorchTPU提供了TensorParallelConfig：在from_pretrained阶段就把transformer切分到多颗芯片上加载，完整模型无需装进单卡。多数常用pipeline已开箱支持，其余只差一份tp_plan。

官方给出的实测基于Cloud TPU v6e-4，四路张量并行、bf16、batch size为1。视频与音频生成的MiniMax-H3在1344×768、5.2秒、49步设置下，编译模式耗时188秒，eager模式为262秒；图像模型Qwen-Image-2.1在1024×1024、40步下为3.5秒对8.0秒，2048×2048的示例同样跑通。

这些数字说明的是编译带来的稳定收益，而不是TPU全面胜过GPU。真正的结构性变化在于抽象层：TPU从“需要移植的目标”变成“另一个device”，PyTorch生态里的模型定义、pipeline与内核可以不经分叉地落到另一种加速器上。

对基础设施团队而言，这降低了多硬件供应的迁移成本，也让TPU在推理与小规模微调场景中更容易被评估。接下来的观察点是更大规模分布式训练的实际表现、Kernels Hub的覆盖广度，以及生产环境中图形编译与动态输入带来的运维复杂度。