PipelineRL：透過即時權重更新優化大型語言模型強化學習

PipelineRL 是一個實驗性的強化學習（RL）實作，旨在解決大規模 LLM 訓練中高推論吞吐量與在策略資料收集之間的根本權衡。透過實作 inflight weight updates，PipelineRL 讓推論伺服器在不停止推論流程的情況下接收更新的模型權重，確保 GPU 高利用率，同時保持訓練資料接近在策略。

解決推論吞吐量與在策略之間的權衡

在傳統的 RL 工作流程中，效率與資料新鮮度之間存在衝突。為了達到高吞吐量，推論伺服器通常使用大批次大小，這會產生供多個策略優化步驟使用的資料。然而，每一次後續的優化步驟都會增加用於收集資料的權重與當前策略權重之間的「延遲」，使得資料越來越偏離在策略，對訓練的效用降低。

PipelineRL 透過在每一次 optimizer 步驟後更新推論伺服器的權重，且不停止推論，來緩解此問題。系統僅在接收新權重所需的短暫時間內暫停推論伺服器。此方法讓推論伺服器得以維持最佳批次大小，同時確保資料保持在策略或接近在策略，從而實現更穩定且有效的學習。

效能與穩定性結果

使用在 Open-Reasoner-Zero 資料集上訓練的 7B 與 32B 模型的實驗顯示，PipelineRL 在 AIME 2024 與 MATH 500 推理基準測試上，表現與 Open-Reasoner 相當或更佳。

簡化的 RL 演算法

儘管表現具競爭力，PipelineRL 使用的 RL 實作遠比 Open-Reasoner-Zero 簡單。主要的簡化包括：

  • Simplified GRPO：它使用簡化版的 Group Relative Policy Optimization（GRPO），且不含價值函數。
  • No Complex Filtering：此實作省略了信任區域重要性權重限制、過長序列過濾以及獎勵塑形。
  • Basic Loss Normalization：損失以批次中序列數量作為分母進行正規化，對所有 token 給予相同權重。
  • No Penalties：系統不使用 KL 懲罰或熵獎勵（雖然支援參考模型 KL）。

KV 快取過期資料的影響

即時權重更新的一個主要顧慮是，序列生成會使用 KV 快取中過期的鍵和值，因為它們是以先前版本的模型計算的。然而，實驗結果顯示此情況並未對訓練穩定性產生負面影響。

模組化架構與技術合約

PipelineRL 設計為模組化，以便與專門的推論（例如 SGLang、vLLM）與訓練（例如 DeepSpeed、FSDP、TorchTitan）軟體整合。這透過兩個主要合約實現：

推論合約

要與 PipelineRL 整合，推論軟體必須提供三個特定 API：

  1. Process group initialization：使用 HTTP POST /init_process_group 請求來初始化用於權重更新的 process group。
  2. Weight Update Trigger：使用 HTTP POST /request_weight_update 請求，向推論伺服器發出暫停並透過 NCCL 接收權重廣播的訊號。
  3. Chat completion：標準的 HTTP POST /v1/chat/completion 請求，用於演員互動。

訓練合約

訓練軟體必須提供以下操作的 Python API：

  • Worker initialization：載入並分割訓練權重與 optimizer 狀態。
  • Forward pass：產生 token 的對數似然。
  • Backward step：計算並累積 RL 目標的梯度。
  • Optimizer Step：執行 optimizer 步驟。
  • Weight gathering and broadcasting：逐層收集更新的權重，並廣播至推論伺服器。

實驗配置

PipelineRL 在 7B 與 32B 模型上進行測試，使用以下超參數：

  • Batch size：4096
  • Learning rate：1e-6
  • Max generated tokens：8192

訓練計算需求大約為 7B 模型在 2 台節點上 3.5 天，32B 模型在 4 台節點上 6 天。

Sources