TPU 成 PyTorch 原生裝置 JAX 先至唔使學
Google 同 Meta 喺 2026 年 10 月 20 日聯合開源咗 TorchTPU，一個原生嘅 PyTorch（Python 機械學習框架）後端，專門運行喺 Google Cloud 嘅 TPU（Tensor Processing Unit，張量處理單元）上。呢個直接改寫咗過去幾年 PyTorch 用家想玩 TPU 都有幾癎嘅現狀。
TPU 係 Google 自家嘅 AI 加速器，Google 設計咗大約十年，自己訓練同服務自家模型（包括 Gemini）都用佢。佢喺 Google Cloud 由單片到數千片嘅 pod 都有，最適合大量算力嘅訓練同大規模推理。但一直以來嘅問題係軟件：TPU 嘅堆棧圍繞 TensorFlow 同 JAX 長大，PyTorch 用家想入場都要多走一步。
點解要換
過去喺 TPU 跑 PyTorch，要行 PyTorch/XLA（torch_xla）。佢可以行，但有自己嘅怪癖：tensor 係「lazy」嘅（tensor 都唔係你），你以為個 code 行咗，其實未必；細細個改動都可能觸發圖重編譯，除錯好難諗。Hugging Face 自己就為咗呢個問題先造咗 Optimum TPU，俾 HF 模型喺 TPU 上做推理同訓練。
TorchTPU 一出，呢套繞路就唔使行。預設用 eager 模式（即寫即行，唔會「lazy」），想更快就加 torch.compile 同 XLA 編譯，對新工作負載係取代 PyTorch/XLA 嘅標準選擇。
點樣用、同對 Hugging Face 嘅影響
裝喺 TPU VM 一條指令就得：
pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu
個 extra index 係俾你 CPU build 嘅 torch，唔係個大好多嘅 CUDA build（TPU 用唔上）。TorchTPU 仲包含為 TPU 調優嘅 Pallas 核心，支援分散式訓練。
對 Hugging Face 講，影響直接：Transformers、Diffusers 同 Kernels 而家原生喺 TPU 上運行。揀個模型出 Hub，send 到 "tpu"，完成。Transformers v5 嘅設計就係做模型定義嘅 source of truth，「anywhere」而家包括 TPU；load 模型入 TPU 就 pass device_map="tpu"。Diffusers 入面 TPU 而家係 first-class device，原型除錯用 eager，部署加一個 compile call，code 同 pipeline 一樣，唔使維護 TPU fork。Optimum TPU 就唔再需要。
模姐驗嚇貨
呢個係生態位嘅事，唔係晶片嘅事。Google 有 TPU 十年，但以前想喺 TPU 用 PyTorch，你要有耐性行 XLA，或者轉 JAX；而家 PyTorch 用家可以好似用 CUDA GPU 咁，即寫即行。呢個對 PyTorch 用家係真錫，對 Google Cloud 嘅吸引力都有幫助——以前個「software catch」正正係攞走。對想喺 TPU 大規模訓練或者部署嘅團隊，呢個係實打實嘅選項，唔係公關。
模姐｜深度調教
