PyTorch代码原生跑上Google TPU

原文要点

Hugging Face 发布 TorchTPU：TPU 成为 PyTorch 原生后端设备，不再必须经过 torch_xla 或 JAX，Optimum TPU 也不再是运行 Transformers/Diffusers 的必需项。在 TPU 虚拟机上用一条 pip 命令安装（需加 CPU 版 PyTorch 的 extra-index-url 避免拉取 CUDA 大包），Transformers 用 device_map="tpu" 加载模型，Diffusers 用 pipe.to("tpu") 即可 eager 推理，另支持 torch.compile 与 TensorParallelConfig 分片。原文称在 v6e-4 上 Qwen-Image 1024×1024 图编译后 3.5 秒、eager 8.0 秒，这些数字是作者实验结果。Transformers 在 TPU 上测试通过率 96.5%，已报 34 个 issue、修复 19 个。

这次做成什么

照着本页在一台 Cloud TPU 虚拟机上把现有 PyTorch/Transformers/Diffusers 代码改到 TPU 上跑通：完成安装、一次文本生成、一次文生图，并理解 eager 与编译模式的差别。练习方案：先用小模型验证环境，再换文中同款管线。

准备什么

• 一个 Google Cloud 账号并创建了 TPU VM（如 v6e 系列），拥有 SSH 登录权限
• Python 3 与 pip 环境，磁盘空间足够放模型权重（十 GB 级）
• Hugging Face 账号令牌（HF_TOKEN），用于下载需要授权的门禁模型
• 已有基础 PyTorch 代码阅读能力，无需 torch_xla 或 JAX 经验

具体怎么做

1. 在 TPU VM 上安装 TorchTPU

   SSH 进入 TPU 虚拟机后执行 pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu。extra-index-url 的作用是装 CPU 版 torch，否则 pip 会默认拉体积大得多的 CUDA 版，在 TPU 上完全用不到。装完用 python -c "import torch; print(torch.device('tpu'))" 确认设备可用。

2. 设置令牌并跑通文本生成

   执行 export HF_TOKEN=你的令牌 后，按下方代码材料新建 demo_llm.py：用 AutoModelForCausalLM.from_pretrained(model_id, device_map="tpu") 直接把模型加载到 TPU，分词后 inputs.to("tpu")，调用 model.generate(max_new_tokens=64) 并解码打印。默认 eager 模式，无需任何 XLA 代码。先用小体量模型验证，再换大模型。

3. 用 Diffusers 在 TPU 上出图

   按材料二新建 demo_sd.py：以 torch.bfloat16 加载 StableDiffusion3Pipeline，pipe.to("tpu") 即可 eager 出图；如需提速，加 pipe.transformer.compile(backend="tpu", fullgraph=True, dynamic=False)。注意首次调用会编译图，之后相同尺寸的调用复用图，所以比较速度时要至少跑两次。

4. 模型放不下单芯片时启用张量并行

   模型超过单芯片显存时，在 from_pretrained 传入 TensorParallelConfig，让 transformer 在加载时按芯片数分片，完整模型不必先装进一处。具体配置写法以 Diffusers 的 TPU 指南为准，本站未实测，参数请在原指南核对。

可直接复制的材料

pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu
export HF_TOKEN=hf_换成你自己的令牌

from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "Qwen/Qwen3.5-9B"
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="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))

import torch
from diffusers import StableDiffusion3Pipeline

pipe = StableDiffusion3Pipeline.from_pretrained(
    "stabilityai/stable-diffusion-3.5-medium",
    dtype=torch.bfloat16,
)
pipe.to("tpu")

# 可选：编译模式，首次调用编译，之后同尺寸复用图
pipe.transformer.compile(backend="tpu", fullgraph=True, dynamic=False)

image = pipe("a corgi astronaut planting a flag on a pizza moon").images[0]
image.save("out.png")

输入与输出示例

本站虚构演示，未实际运行；耗时不保证，原文的 3.5 秒/8.0 秒是作者在特定芯片与模型上的结果。

输入

在 v6e-4 TPU VM 上运行 demo_sd.py，提示词：a small robot watering a tiny bonsai on a windowsill, soft morning light，尺寸 1024×1024、bf16。

对应输出

首次调用（含编译）约数十秒后 out.png 落盘，内容为窗台上一台小机器人给盆景浇水的插画；随后以完全相同的尺寸与提示词再调用一次，第二次明显快于第一次，因为编译后的图被复用。若换成 768×768 再跑，会再次经历一次编译。

做完检查这些

• python 中能成功构造 torch.device('tpu') 且导入 torch 不报错
• demo_llm.py 打印出一段以输入开头、继续生成的英文文本，而不是设备相关异常
• demo_sd.py 生成 out.png 且文件可打开，连续两次同尺寸调用第二次更快
• 全过程代码中没有 import torch_xla 或任何 JAX 依赖

容易卡住的地方

• 漏掉 --extra-index-url 会安装 CUDA 版 torch，白白占用大量磁盘且不需要
• 编译模式首次调用慢是正常现象，拿首次耗时判断性能会得出错误结论
• 张量并行分片并非所有管线都开箱支持，不支持的模型需要自行定义 tp_plan
• 原文称 TPU 测试通过率 96.5%，仍有失败项，遇到异常应先到 TorchTPU 仓库查 issue 而不是怀疑自己的代码
• 门禁模型未在 Hub 接受条款或未设 HF_TOKEN 时，下载阶段就会失败