Hugging Face 介绍 TorchTPU：为 PyTorch 增加原生 TPU 后端，Transformers 与 Diffusers 可直接调用

OmniTools 10月20日消息，Hugging Face 10月20日发布博客介绍 TorchTPU，这是 PyTorch 的原生 TPU 后端，使 TPU 成为 PyTorch 中的一个普通设备。此前在 TPU 上运行 PyTorch 需借助 torch_xla 或转向 JAX，Hugging Face 也曾推出 Optimum TPU 提供推理与训练支持；官方表示，在使用 Transformers 和 Diffusers 时，TorchTPU 已不再需要 Optimum TPU。

在具体用法上，Transformers 可通过 device_map="tpu" 将模型直接加载到 TPU，默认以 eager 模式运行，也可结合 torch.compile 或分布式配置；Diffusers 中只需 pipe.to("tpu")，编译模式额外增加一次 compile 调用。对于单芯片放不下的模型，可传入 TensorParallelConfig，在加载时按芯片切分。博客还给出安装命令 pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu。

TorchTPU 附带针对 TPU 调优的 Pallas 内核、分布式训练支持，以及与 torch.compile 的集成。Kernels Hub 提供可直接取用的 TPU 内核，示例中将 Tokamax 的 ragged-dot 内核接入 Qwen3-30B-A3B 等 MoE 模型，批大小为 64 时解码分别提速 21.5 倍和 12.8 倍（对比 TPU v6e-8 上的默认实现）。

测试方面，官方称在 TPU 上运行的 Transformers 测试通过率为 96.5%，目前已向 TorchTPU 仓库提交 34 个问题，其中 19 个已修复；TPU 测试还发现 StaticCache 与 Qwen3 张量并行计划等 Transformers 自身缺陷，并已在上游修复。团队表示正推进 TPU 上的夜间 CI，并将结果展示在 Transformers CI 面板中。