返回首頁

AI 基礎設施與訓練系統

LazyTrain 用混合整數排程重疊記憶體搬移,單張 H800 訓練 27B 模型快 1.24 倍

LazyTrain 在訓練前聯合決定 activation 保存、卸載及重算位置,把 PCIe 與 NVMe 傳輸塞入 GPU 計算窗口。Qwen3.6‑27B 實驗由 176.90 提高至 219.95 TFLOPS,但目前只驗證單卡與離線排程。

Adam-dalekie-pole · CC BY-SA 3.0 · Image source
zh-Hant

記憶體不足時,單卡大型模型訓練通常在 activation checkpointing、CPU offload 與重算之間採用固定規則。LazyTrain 把這些選擇改寫成混合整數線性規劃:輸入 GPU、主記憶體與 NVMe 容量,以及 PCIe、SSD 頻寬和各層計算時間,再一次求出哪些邊界需要保存、放在哪一層儲存體、哪些區塊重算,以及傳輸應落在哪個計算窗口。

系統建立從反向傳播起點到終點的路徑流約束,同時限制 HBM、DRAM、NVMe 容量及參數、梯度與 activation 的雙向流量。SCIP 求解器在訓練前產生排程,runtime 執行時不再即時搜尋。主要 Qwen3.6‑27B 設定有 64 個 Transformer block、序列長度 1,024、batch 72;52 個實體化 checkpoint 中,30 個留在 GPU、21 個置於 CPU、1 個落到 NVMe,參數及 optimizer state 則以 CPU 為主。

在同一張 H800、相同資料切分與 batch 下,LazyTrain 相對 MegaTrain 將持續運算量由 176.90 提高至 219.95 TFLOPS,約為 1.24 倍,吞吐達每秒 1,361 token,GPU 峰值約 68.84GB。移除混合整數排程後,效能降至 193.17 TFLOPS,顯示主要增益來自傳輸與重算的聯合配置;所謂 Hybrid 8-bit 路徑只令移除後的數值降至 219.29 TFLOPS,而且量化狀態僅用於約 2% 參數。

開源 alpha 建基於 MegaTrain 的 CPU-master、逐層串流架構,並提供 Qwen3.6‑27B、MetaMathQA 與 NVMe 排程設定。不過論文只做單卡 H800 及 RTX 3090 實驗,每列通常僅一次執行;排程不會因執行期間的頻寬抖動重新配置,宣稱的零額外曝露傳輸時間也是求解器估值,未以逐步 stall instrumentation 驗證。更值得關注的是多 GPU、跨節點及不同 MoE 架構能否維持增益,而非單一 27B 案例的峰值。

來源

  1. LazyTrain: Limited-resource Allocation toward Zero-waste Yield Optimization in Large Language Model Training
  2. DataArcTech/LazyTrain