Hugging Face TRL：共同定位的 vLLM 以提升 GRPO 訓練效率

Hugging Face 已在 TRL 函式庫中加入共同定位 vLLM 的支援，使訓練與推論能共享相同的 GPU。這消除了伺服器模式推論的「乒乓」低效情況，訓練 GPU 在生成時閒置、推論 GPU 在訓練時閒置，從而提升吞吐量並降低硬體需求。

伺服器模式推論的低效

在 TRL v0.18.0 之前，vLLM 僅支援伺服器模式。在此配置下，vLLM 以獨立進程在專用 GPU 上執行，透過 HTTP 與訓練腳本通訊。此架構在像群組相對策略最佳化 (GRPO) 這類線上學習演算法中會產生顯著瓶頸，因為生成持續不斷：

  • GPU 利用率不足：訓練 GPU 在 vLLM 伺服器生成完成時閒置。
  • 資源浪費：必須專門配置僅供推論使用的 GPU，增加成本。
  • 通訊開銷：依賴 REST API 呼叫會帶來網路延遲。

共同定位的設計與實作

共同定位允許 vLLM 與訓練程式碼在同一個分散式進程組中同時執行。訓練器不再使用外部伺服器，而是透過外部啟動器在進程內啟動 vLLM。

主要技術能力

  • 統一執行：訓練與推論任務交替使用相同的 GPU 資源，減少閒置時間。
  • 直接記憶體通訊：移除 HTTP 呼叫後，vLLM 透過原生 Python 呼叫與訓練迴圈溝通。
  • 分散式相容性：系統相容 torchrun、張量平行 (TP) 與資料平行 (DP)，採用單程式多資料 (SPMD) 執行模式。
  • 簡化部署：vLLM 直接在訓練作業內受控，無需額外的伺服器腳本。

設定

使用者可透過在 GRPOConfig 中設定 vllm_mode="colocate" 來啟用此功能：

training_args = GRPOConfig(
    ...,
    use_vllm=True,
    vllm_mode="colocate",
)

效能基準測試

實驗比較伺服器模式（8 顆 GPU 中有 1 顆專屬於 vLLM）與共同定位模式（全部 8 顆 GPU 用於訓練），顯示出一致的吞吐量提升。為確保公平，伺服器模式的吞吐量以 8/7 的係數進行正規化。

各模型結果

  • 1.5B Model：1.5B 模型：在最大批次大小下達到最高 1.43 倍的加速。然而，提升張量平行 (TP) 會因通訊開銷而降低效能。
  • 7B Model：7B 模型：在不同批次大小下達到最高 1.35 倍的加速，提升 TP 時最高可達 1.73 倍，顯示較大的模型從分片中受益更多。
  • 72B Model (Qwen2.5-Math-72B)：72B 模型（Qwen2.5-Math-72B）：共同定位設定比純 TRL 快約 1.26 倍，即使使用的 GPU 數量少 4 顆。

擴展至大型模型（72B 以上）

訓練如 Qwen2.5-Math-72B 這樣的大型模型需要先進的記憶體管理，以避免訓練與生成之間的資源競爭。

vLLM 睡眠模式

Hugging Face 將 vLLM 的 sleep() API 整合至 GRPO 迴圈。使用 Level 2 睡眠將模型權重與 KV 快取全部從 GPU 卸載，最大化訓練步驟的可用記憶體，避免記憶體競爭。

DeepSpeed 與 Accelerate 優化

為了維持穩定性與效率，採用了以下優化措施：

  • DeepSpeed ZeRO Stage 3：DeepSpeed ZeRO Stage 3：將模型權重、梯度與優化器狀態在 GPU 之間切分。
  • CPU 離線：CPU 離線：將優化器狀態移至 CPU ("offload_optimizer": {"device": "cpu"}) 以釋放 GPU 記憶體。
  • 記憶體管理：記憶體管理：啟用 contiguous_gradients 以減少記憶體碎片，並使用 overlap_comm 加速訓練。
  • Accelerate 整合：Accelerate 整合：用於多 GPU/多節點啟動與資料平行。

挑戰與經驗教訓

共同定位 vLLM 的實作揭露了多項技術挑戰：

  • vLLM ≥ 0.8.0 的確定性：外部啟動器的張量平行錯誤被追溯到需要明確設定隨機種子。
  • 緩衝區還原：Level 2 睡眠的錯誤（Issue #16564）需要修正，以在喚醒後明確還原模型緩衝區（例如 BatchNorm 的 running mean/var）。
  • 關閉穩定性：已知問題是 vLLM 睡眠在結束訓練過程時會導致 segmentation fault（Issue #16993），但不影響訓練結果。

結論：模型品質

在 Math500 與 AIME24 的基準測試證實，共同定位不會削弱模型效能。共同定位與純 TRL 設定的獎勵曲線幾乎相同，且共同定位訓練的模型表現與純訓練模型相當，驗證了效率提升並未以準確度為代價。

Sources