Hugging Face 推出 TorchTPU，為 PyTorch 加入原生 TPU 支援

如果你曾經試過在 TPU 上執行 PyTorch，可能曾把程式碼移植到 torch_xla，或者最後被迫學習 JAX。有了 TorchTPU 這個原生 PyTorch 後端，TPU 現在就只是另一種裝置。它讓 TPU 在 Hugging Face 生態系統中獲得原生支援。Transformers、Diffusers 和 Kernels 都可在 TPU 上原生執行。從 Hub 選擇一個模型，將它傳送到 "tpu"，就完成了。

等等，TPU 是甚麼？🤔

TPU 是 Tensor Processing Unit 的縮寫。TPU 是 Google 自家的 AI 加速器。這些晶片專為一項工作而設：神經網絡大部分時間都在進行的矩陣運算。Google 設計 TPU 約有十年，並用它們訓練及提供自家模型的服務，包括 Gemini。

TPU 可透過 Google Cloud 使用，規模由單一晶片到數以千計晶片組成的 Pod 不等。無論是訓練，還是大規模提供服務，只要需要大量運算，TPU 就能發揮所長。

一直以來，問題都在於軟件。TPU 軟件堆疊是圍繞 TensorFlow 和 JAX 發展的，PyTorch 用戶要花更多工夫才能使用 TPU。

以往的做法

直到現在，要在 TPU 上執行 PyTorch，都必須透過 PyTorch/XLA（torch_xla）。它雖然可用，但有一些特性需要適應。張量採用惰性求值（惰性的是張量，不是你），所以程式碼不一定會在你預期的時候執行。小幅修改可能觸發計算圖重新編譯，而除錯時也很難掌握執行過程。

為了讓使用體驗更簡單，我們開發了 Optimum TPU，讓 Hugging Face 模型可在 TPU 上進行推理和訓練。現在 TorchTPU 提供原生支援，使用 Transformers 和 Diffusers 在 TPU 上執行 Hugging Face 模型已不再需要 Optimum TPU。

TorchTPU 登場 🚀

TorchTPU 為 PyTorch 加入 TPU 作為新後端。你只需編寫一般的 PyTorch 程式碼。它配備針對 TPU 調校的 Pallas kernels、分散式訓練支援，以及 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 mode 執行。你也可以使用 torch.compile 以提升效能，或在分散式環境中執行。

圖像？影片？Diffusers 在 TPU 上繪製它們 🎨

TPU 現在是 Diffusers 的一級裝置。進行原型開發或除錯時，eager mode 只需 pipe.to("tpu")。部署時，編譯模式只需多加一次 compile 呼叫。你的 pipelines 和程式碼都毋須更改，也不必維護 TPU 專用分支。

import torch
from diffusers import StableDiffusion3Pipeline

pipe = StableDiffusion3Pipeline.from_pretrained(
    "stabilityai/stable-diffusion-3.5-medium", dtype=torch.bfloat16
)
pipe.to("tpu")  # eager mode: that's the whole change
# optional: compiled mode
pipe.transformer.compile(backend="tpu", fullgraph=True, dynamic=False)

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

第一次呼叫會進行編譯；之後只要圖像尺寸相同，每次呼叫都會重用計算圖。

模型太大，無法放進單一晶片？將 TensorParallelConfig 傳入 from_pretrained，載入時就會把 transformer 分片至多個晶片，因此無須將完整模型放進單一晶片。許多熱門 pipelines 開箱即用便支援這項功能；尚未支援的，只需加入一個 tp_plan 即可。Diffusers TPU 指南提供完整做法。

以下是一些在 Cloud TPU v6e-4 上執行的示例：

_{TP=4（每個模型分拆至四個晶片）· bf16 · batch 1}

MiniMax-H3（影片 + 音訊）	Qwen-Image-2.1
"圖片中的毛茸茸 Hugging Face 吉祥物玩偶，一張圓圓的黃色笑臉，面頰紅潤，還有兩隻張開的小手，站在一艘古希臘三列槳戰船的船首 […] 海浪拍打聲、木材吱嘎聲、帆索間的風聲、遠處的海鷗聲，以及低沉而磅礡的管絃樂漸強聲。" · 1344×768, 5.2 s, 49 步 · 編譯模式 188 s，急切模式 262 s	"一隻戴着巫師帽的水豚，在燭光下看書，油畫" · 1024×1024, 40 步 · 編譯模式 3.5 s，急切模式 8.0 s	"一個霓虹燈店舖招牌，上面寫着 'QWEN IMAGE 2.1'，雨夜，濕漉漉的路面倒映着燈光" · 2048×2048, 40 步

核心算子

Kernels Hub 提供可直接使用的 TPU 核心算子。使用 Kernels library 載入其中一個，即可提升模型速度。

TPU 核心算子以 Pallas 編寫，讓你能以低階方式控制硬件。你也可以使用 Tokamax library 的核心算子。以下是將 Tokamax 的 ragged-dot 核心算子接入混合專家模型的例子：

from kernels import get_kernel
from transformers import AutoModelForCausalLM
from transformers.integrations.moe import ALL_EXPERTS_FUNCTIONS

repo_id = "tengomucho/ragged-dot-tpu"
kernel = get_kernel(repo_id, version=1, trust_remote_code=[repo_id])
ALL_EXPERTS_FUNCTIONS.register("tokamax_ragged_dot", kernel.experts_forward)
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-30B-A3B-Instruct-2507",
    experts_implementation="tokamax_ragged_dot",
)

批次越大，效能提升越明顯。批次大小為 64 時，Qwen3-30B-A3B 的解碼速度快 21.5×，Gemma 4 26B-A4B 則快 12.8×。

在 TPU v6e-8 上，Hub 核心算子與預設實作的解碼吞吐量比較。

我們替你測試 🧪

我們持續透過執行程式庫測試套件，改善 transformers 對 TPU 的支援。我們正努力加入 TPU 夜間 CI，並在 (Transformers 儀錶板)[https://huggingface.co/spaces/transformers-community/transformers-ci-dashboard] 公開測試結果。

Transformers CI 儀錶板，顯示 TPU 與其他加速器的測試結果。仍在開發中！

測試結果尚未全部亮起綠燈，我們可以接受。在 TPU 上執行的測試中，96.5% 通過。每次發現測試失敗，我們都會編寫一個小型重現腳本，並在 TorchTPU 儲存庫提交問題回報。目前我們已回報 34 個問題，其中 19 個已修復。在 TPU 上執行測試，也找出 Transformers 本身的一些錯誤，例如 StaticCache 和 Qwen3 張量平行計劃中的錯誤。我們已在上游修正這些問題。

如果 TPU 上出現問題，我們希望在你發現之前先找出來。

開始使用

對 PyTorch 用戶而言，TPU 不再是支線任務。安裝 TorchTPU，從 Hub 取得模型，然後將模型傳送到 "tpu"。

• GitHub 上的 TorchTPU
• Hub 上的 TPU 核心算子
• Cloud TPU 文件

我們很期待看看你會打造甚麼 🤗。如果遇到問題，請提交 issue 告訴我們。