Google MaxText 導入彈性訓練:TPU 節點故障後幾秒內自動復原
Google Developers Blog · 2026-07-06
Google 工程團隊在開源訓練框架 MaxText 上做了一次刻意的故障注入實驗:訓練途中主動砍掉一整個 TPU slice,系統從偵測到故障、重新排程到恢復訓練只花了約 110 秒,僅損失 88 個訓練步。這項能力稱為彈性訓練(elastic training),核心是把 Pathways 與 Orbax checkpoint 機制結合,讓分散式訓練不必因單點硬體故障而整個作業重啟。
背景
大規模 TPU 訓練慣例採用 SPMD 架構,每個運算節點各跑一份相同的 Python 程序;一旦某台機器的 TPU 掉線,整個作業通常得從最近一次 checkpoint 完整重啟,重新排程、重新載入權重、重建通訊拓撲。MaxText 改採單一控制器架構(single controller),讓跑在 CPU 節點(n2-standard-64)上的單一 Python 程序,透過 Pathways proxy 把所有 TPU 晶片都視為本地裝置。這樣一來,某個 TPU slice 掛掉時,健康的控制器程序仍然存活,能直接攔截例外並啟動復原邏輯,不需要整個 JobSet 重新排程。
核心改動:兩種復原路徑
MaxText 提供兩種彈性復原策略。Pause and resume 會等待故障 slice 的替補資源就緒,再把 checkpoint 重新載入到完整運算網格恢復訓練;replica resize 則讓訓練在剩餘存活的 slice 上以較低吞吐量繼續,同時背景等待替補節點加入。兩者依賴同一套故障偵測與 checkpoint 驗證機制:
- Pathways 透過
DATA_LOSS錯誤偵測進行中操作的失敗,或在約 10 秒未收到心跳後拋出DEADLINE_EXCEEDED - 上述錯誤會包裝成可攔截的
jax.errors.JaxRuntimeError,交給pathwaysutils提供的elastic_retrydecorator 處理 elastic_retry攔截例外、記錄事件,並在同一程序內執行復原,不讓程序崩潰- Orbax 驗證 checkpoint 完整性:檢查目錄下是否有
commit_success標記檔;若因寫入中途失敗而缺失,該目錄會被刪除,復原改用最新的合法 checkpoint
啟用彈性訓練只需在 MaxText 訓練指令加入對應旗標:
enable_single_controller=true
elastic_enabled=true
elastic_timeout_seconds=300
elastic_max_retries=10
checkpoint_period=100提交工作負載時,xpk 工具也要加上對應設定,例如 --elastic-slices=3、--max-slice-restarts=10,JobSet 的 backoffLimit 設為 20,允許 slice 層級的 pod 重啟而不觸發整個 JobSet 重建。若模型規模更大,官方建議加上 ENABLE_PATHWAYS_PERSISTENCE=1,讓 TPU host 直接讀寫 Cloud Storage 上的 checkpoint 分片,避免流量都擠在 Pathways proxy 造成記憶體瓶頸。
實驗結果
示範實驗使用 Qwen3-0.6B 模型,硬體配置為 3 個 TPU v5e-16 slice(共 48 顆晶片)加一個 CPU 控制器節點,JobSet 包含 1 個 head pod 與 12 個 worker pod(每 slice 4 個 host)。研究人員在訓練中途手動終止其中一個 TPU slice 後記錄到以下耗時:
| 階段 | 耗時 |
|---|---|
| 心跳逾時到拋出錯誤 | 約 13 秒 |
| Pod 重新排程與歸隊 | 約 50 秒 |
| Checkpoint 復原(7 GiB) | 5.39 秒 |
| 第一步 warmup | 12.7 秒 |
| 總停機時間 | 約 110 秒 |
整個流程中,控制器程序的 PID 與 Python 物件狀態始終保留在 CPU 記憶體裡,已編譯的 XLA cache 也留在 Cloud Storage 上不必重新編譯,真正損失的只有最近一次 checkpoint commit 之後的 88 個訓練步。
訓練政策進步不等於推論政策進步:天津大學與阿里提出 LLM 強化學習新目標 MIPU
arXiv:2606.29526 · 2026-07-06
天津大學與阿里巴巴的研究團隊發表論文,指出當前 LLM 強化學習普遍存在一個被忽視的目標錯位:演算法優化的是訓練引擎裡的策略,但真正部署上線的是推論引擎產生的策略,兩者在存在「訓練—推論不一致」(training-inference mismatch)時會分道揚鑣。論文提出 Monotonic Inference Policy Improvement(MIPI)作為新的優化原則,並實作對應演算法 Monotonic Inference Policy Update(MIPU),論文編號 arXiv:2606.29526。
背景
在大規模 LLM RL 訓練中,推論(生成 rollout)與訓練(計算梯度)通常跑在不同的引擎上,即便兩邊參數已經同步,浮點運算路徑、算子實作、量化策略等差異仍會讓同一組參數在推論引擎與訓練引擎算出不同的 token 機率。這種不一致並非新問題,但過去的做法幾乎都聚焦在如何讓「訓練側」的策略更新更穩定,論文團隊認為這個假設本身就有漏洞:訓練政策的改善,在不一致存在時,並不保證推論政策(也就是實際部署行為)也跟著改善。
方法細節
MIPI 原則把優化目標重新對齊到部署端,要求每一次策略更新都必須是對推論政策的單調改善,而不只是對訓練政策的改善。基於這個原則,論文設計了兩步驟框架 MIPU:
- 第一步:以取樣器(sampler)為參照,構造候選的策略更新,產生一組「sampler-referenced」候選解
- 第二步:同步候選解之後,透過一個推論側差距代理指標(inference-gap-aware acceptance filtering)來篩選——只有能通過推論側落差檢驗的候選更新才會被接受並套用
這種設計等於在每次更新前後都插入一道以推論政策表現為準的驗收關卡,把「訓練引擎算出來政策變好了」與「推論引擎實際部署也變好了」這兩件事重新綁在一起,避免優化過程只是在訓練側自欺欺人地進步。
實驗結果
論文在兩種模型規模、高度不一致的訓練—推論條件下測試 MIPU,結果顯示相較於只針對訓練政策做優化的做法,MIPU 在平均推理表現與訓練穩定性上都有提升。論文未在摘要與初版頁面公開具體的基準資料集、比較基線演算法與量化數字,完整實驗設置需參照論文全文與附錄。
LeRobot v0.6.0 發布:世界模型 policy、獎勵模型,把失敗案例變回訓練資料
Hugging Face Blog · 2026-07-07
Hugging Face 機器人學習框架 LeRobot 釋出 v0.6.0,主題是把機器人學習迴圈補完整:讓 policy 在動作前先「想像」未來、用獎勵模型判斷任務是否成功、再用新的部署 CLI 把失敗案例直接轉成下一輪訓練資料。這個版本一次新增三種世界模型 policy、五個 VLA 模型整合、六個模擬基準,以及重寫過的資料載入與部署流程。
核心改動:世界模型 policy 與獎勵模型
三種新的世界模型 policy 各自用不同方式把「預測未來」接進動作生成:
VLA-JEPA:以Qwen3-VL-2B為基礎的精簡 VLA,訓練時在潛在空間預測未來影格,但世界模型本身在推論階段會被丟棄,只留下訓練時學到的監督訊號,不增加推論成本LingBot-VA:自迴歸的影片—動作模型,同時預測未來影片與動作,並以真實觀測為條件;支援--policy.save_predicted_video=true存下預測影片,單張 24–32 GB GPU 即可運行FastWAM:搭配約 5B 參數的影片生成專家與精簡動作專家,跳過測試時的「做夢」步驟,直接對動作區塊(action chunk)做去噪
獎勵模型方面新增 Robometer,以 Qwen3-VL-4B 為底、在超過 100 萬條機器人軌跡上訓練,能根據影片與語言指令直接評分任務進度與成敗;另一個 TOPReward 則是零樣本做法,直接讀取任意夠強的 VLM 在給定軌跡與指令下輸出「True」token 的對數機率作為成功分數,不需要額外訓練。
影響範圍:VLA 模型整合與部署工具鏈
此版本也整合了多個第三方 VLA 模型,包括 NVIDIA GR00T N1.7(改用 Cosmos-Reason2-2B VLM 與 flow-matching 動作頭,並與 Isaac-GR00T 做過對照測試)、Allen Institute for AI 的 MolmoAct2(bf16 推論約需 12 GB 顯存)、Qwen2.5-VL-3B 骨幹的 EO-1,以及參數量分別約 450M 與 0.77B 的 Multitask DiT、EVO1。評測端新增六個模擬基準,涵蓋 LIBERO-plus(一萬組 LIBERO 擾動變體)、RoboTwin 2.0(50 個雙臂任務、逾 10 萬條軌跡)、RoboCasa365(2500 個程序生成廚房場景、365 項任務)等,統一透過以下指令執行:
lerobot-eval --policy.path=... --env.type=... --eval.n_episodes=100部署面新增 lerobot-rollout CLI,支援 base、sentry、highlight、episodic 與 dagger 五種部署策略,其中 dagger 是 DAgger 風格的人機協同修正,會替人類介入的片段打上標籤,直接產生可回收利用的訓練資料。資料載入端則導入 NVENC、VideoToolbox、VAAPI、QSV 等硬體加速影片編碼探測、12-bit 深度串流支援,以及可讀取子集合、速度提升兩倍的資料載入器,子集合載入時間從 275 秒降到 0.06 秒。訓練端支援透過 Accelerate 做 FSDP 全分片資料平行,並可用 --job.target=a10g-small 提交到 HF Jobs,涵蓋 T4 到 8x H200 各級硬體。
原始來源:Hugging Face Blog