Hugging Face ❤️ TorchTPU：TPU 原生支持

如果你曾尝试在 TPU 上运行 PyTorch，可能不得不把代码移植到 torch_xla，或者最终去学 JAX。 有了 TorchTPU 这个原生 PyTorch 后端，TPU 现在只是另一种设备。 它为 Hugging Face 生态带来了对 TPU 的原生支持。

Transformers、Diffusers 和 Kernels 都能在 TPU 上原生运行。 从 Hub 选一个模型，把它发送到 "tpu"，就完成了。 等等，TPU 是什么？ 🤔 TPU 指张量处理单元（Tensor Processing Unit）。

TPU 是谷歌自研的 AI 加速器。 这些芯片专为一件事打造：神经网络大部分时间都在做的矩阵运算。 谷歌设计它们已有约十年，并用它们训练和部署自家模型，包括 Gemini。 TPU 可在 Google Cloud 上使用，从单芯片到数千芯片的 pod 都有。

当你需要大量算力时——无论是训练还是大规模推理服务——它们表现出色。 问题一直出在软件上。 TPU 软件栈是围绕 TensorFlow 和 JAX 成长起来的，PyTorch 用户想接入要费更多功夫。 旧方式 到目前为止，在 TPU 上运行 PyTorch 意味着要经过 PyTorch/XLA（torch_xla）。

它能用，但有怪癖。 张量是惰性的（说的是张量，不是你），所以代码并不总在你以为的时候执行。 小改动可能触发图重编译，调试也很难推理。 为了让生活更轻松，我们构建了 Optimum TPU。 它为 Hugging Face 模型提供了在 TPU 上的推理和训练支持。

现在 TorchTPU 带来了原生支持，用 Transformers 和 Diffusers 运行 Hugging Face 模型不再需要 Optimum TPU。 TorchTPU 登场 🚀 TorchTPU 为 PyTorch 增加了 TPU 这一新后端。

你写常规 PyTorch 代码。 它附带为 TPU 调优的 Pallas 内核、分布式训练支持，以及在你想要更高速度时的 torch.compile 集成。 在 TPU VM 上安装只需一条命令： pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu 额外索引提供 torch 的 CPU 构建。

没有它，pip 会拉取大得多的 CUDA 构建，而你在 TPU 上并不需要。 你喜欢的库，现在都能在 TPU 上运行 Transformers Transformers v5 围绕一个理念构建：成为模型定义的权威来源，让它们能在任何地方运行。

"任何地方"现在包括 TPU。 要在 TPU 上加载模型，传入 device_map="tpu"。 就这样。 from transformers import AutoModelForCausalLM, AutoTokenizer model_id = "Qwen/Qwen3.5-9B" model = AutoModelForCausalLM.from_pretrained( model_id, device_map="tpu", # will load model directly on the TPU ) tokenizer = AutoTokenizer.from_pretrained(model_id) inputs = tokenizer("The future of AI is", return_tensors="pt").to("tpu") out = model.generate(**inputs, max_new_tokens=64) print(tokenizer.decode(out[0], skip_special_tokens=True)) 模型默认以 eager 模式运行。

你也可以用 torch.compile 获得更好性能，或在分布式设置中运行。 图像？ 视频？ Diffusers 在 TPU 上作画 🎨 TPU 现在是 Diffusers 中的一等设备。 对于原型设计或调试，eager 模式只需 pipe.to("tpu")。

对于部署，编译模式增加一次 compile 调用。 你的流水线和代码保持不变，也没有需要维护的 TPU 分支。 import torch from diffusers import StableDiffusion3Pipeline pipe = StableDiffusion3Pipeline.from_pretrained( "stab