降低長上下文混合專家訓練中的記憶體峰值瓶頸
為何重要
長上下文與巨型混合專家模型是目前大型語言模型（LLM）的前沿方向，但訓練門檻極高。這項研究提供了具體的演算法解決方案，在維持精確度的前提下打破記憶體限制。對業界而言，這意味著開發者能在現有硬體規格下進行更大規模模型的訓練或遷移，降低訓練大型模型的物理門檻，並推進中長篇輸出或深度推理能力的實用化。
混合專家模型在長上下文或大批次訓練時，經常因任一元件的記憶體峰值超過裝置容量而失敗。本研究提出四種最佳化排程與計算策略，透過側路、投影與最佳化器層面的調整，來解決每個個別的記憶體峰值瓶頸。
- 識別出四個未受常用並行計畫約束而呈現無界限增長的記憶體峰值：專家分派與路由矩陣、詞彙投影計算、梯度檢查點邊界 (depth times sequence length)，以及最佳化器狀態 (parameter count)。
- 引入四種新方法並在不改變損失與梯度的情況下最佳化運算順序與資料移動：PipelinedLLEP 限制每個分派區塊的 Token 數、Ring-DTP 將 activation 或 weight shard 繞環狀流通並線上性計算 log-sum-exp、Selective checkpoint offload (SCO) 將檢查點的長期張量保留在 CPU 記憶體，以及 OffloadStreamAdamW 改善 CPU 端 Adam 更新的串列流程。
- 在 120B 至 667B 引數的 MoE 模型上進行測試，能支援 100 萬上下文長度，之後兩項指標分別較 tuned FSDP2 基準提升 8–32 倍與高達 10.4 倍吞吐量。