從 torch_xla 到原生後端：簡化 PyTorch 與 TPU 的整合
過去要在 Google TPU 上執行 PyTorch，開發者必須將程式碼移植至 torch_xla，或是學習 JAX。torch_xla 雖然可行，但存在張量惰性執行（lazy execution）的問題，導致程式碼執行時機難以預測，且微小修改可能觸發圖重編譯，增加除錯難度。
Hugging Face 曾推出 Optimum TPU 以協助模型在 TPU 上運行，但隨著 TorchTPU 提供原生支援，Optimum TPU 已不再是執行 Transformers 和 Diffusers 模型的必要條件。TorchTPU 將 TPU 作為 PyTorch 的新後端加入，允許使用者撰寫標準的 PyTorch 程式碼，並包含針對 TPU 調校的 Pallas kernels、分散式訓練支援以及 torch.compile 整合。
在 TPU VM 上安裝 TorchTPU 只需執行一條指令：pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu。此額外索引確保安裝的是 CPU 版本的 torch，避免 pip 自動下載在 TPU 環境中不需要且體積較大的 CUDA 版本。
Transformers 與 Diffusers 的原生 TPU 支援細節
Transformers v5 的設計理念是作為模型定義的單一事實來源，使其能在任何地方運行，現在這也包括 TPU。載入模型至 TPU 時，只需在 from_pretrained 中設定 device_map="tpu"，即可直接將模型載入 TPU 記憶體。預設情況下模型以 eager mode 執行，使用者也可選擇使用 torch.compile 來提升效能或進行分散式設置。
Diffusers 現在也將 TPU 視為一等公民裝置。對於原型開發或除錯，只需呼叫 pipe.to("tpu") 即可在 eager mode 下運行；若需部署，則可額外調用 compile 方法。程式碼與管線無需維護特定的 TPU 分支。首次呼叫會進行編譯，後續相同圖片尺寸的呼叫將重用已編譯的圖形以加速處理。
對於單顆晶片無法容納的大型模型，Diffusers 支援透過 TensorParallelConfig 在載入時將 transformer 分片至多顆晶片上，避免整個模型必須符合單一裝置的限制。許多流行的管線已預設支援此功能，其他管線則可透過 tp_plan 進行配置。
Kernels Hub 提供優化核心以提升推理速度
Kernels Hub 提供了可直接使用的 TPU kernels，這些核心使用 Pallas 語言編寫，允許對硬體進行底層控制。開發者可透過 Kernels 函式庫載入這些核心以加速模型，例如從 Tokamax 函式庫引入 ragged-dot kernel 並註冊到 Mixture-of-Experts (MoE) 模型中。
根據 Hugging Face 公布的測試數據，在 TPU v6e-8 硬體環境下，使用 Kernels Hub 提供的核心相較於預設實作有顯著的速度提升。具體而言，在批次大小為 64 的情況下，Qwen3-30B-A3B 模型的解碼速度提升了 21.5 倍，Gemma 4 26B-A4B 模型則提升了 12.8 倍。這些數據僅適用於特定的硬體配置與批次條件。
測試覆蓋率與已知限制
Hugging Face 持續透過執行 Transformers 函式庫的測試套件來改進 TPU 支援，並計畫在 transformers dashboard 上提供夜間 CI 結果。目前並非所有測試都能通過，但在於 TPU 上執行的測試中，通過率為 96.5%。
開發團隊已報告了 34 個與 TorchTPU 相關的問題，其中 19 個已修復。在 TPU 上的測試也發現並修復了 Transformers 本身的一些錯誤，例如 StaticCache 和 Qwen3 tensor-parallel plans 中的問題。這些修復已合併至上游程式碼庫。
