vLLM V1 移植:確保強化學習中的後端正確性
ServiceNow AI 成功將其 RL 推論引擎從 vLLM V0 移植至 V1,並以後端正確性優先於目標端的修正。團隊透過解決四項關鍵差異,使 vLLM 0.8.5(V0)與 vLLM 0.18.1(V1)達成等價:已處理的 rollout logprob、V1 特有的執行時預設值、進行中的權重更新路徑,以及 lm_head 投影的精度。
移植目標
在像 PipelineRL 這樣的系統中使用 vLLM 作為 rollout 生成的推論引擎時,該引擎會抽樣 token 並回傳 logprob,訓練器利用這些 logprob 計算關鍵的 RL 指標,包括 policy ratio、KL 散度、clip rate 與 entropy。若這些 logprob 的計算方式有任何差異,會產生訓練‑推論不匹配,從而根本改變訓練動態。
為了遷移至 vLLM V1,ServiceNow AI 設定了明確的目標:驗證 V1 回傳的 rollout logprob 符合訓練器的預期格式,對照 V0 參考重新執行工作負載,且僅在後端等價恢復後才評估目標層面的變化。最初使用 V1 的嘗試顯示,獎勵、entropy 與 clip rate 與 V0 參考相比有顯著偏差。
後端失效模式
團隊將可能導致訓練‑推論不匹配的原因分為三層:
- 語意不匹配:後端回傳的 logprob 與訓練器預期的意義不同。
- 推論路徑不匹配:快取、排程或請求處理等執行時預設值的差異,使提示走向不同的執行路徑。
- 目標不匹配:RL 目標需要針對剩餘的陳舊或後端不匹配進行修正。
將前兩項視為後端行為問題並先行排除,團隊避免了過早在目標端套用修正以掩蓋後端錯誤的錯誤做法。
vLLM V1 後端修正
Logprob 語意
預設情況下,vLLM V1 會從原始模型輸出(在 logits 後處理之前,如溫度縮放與 top‑k/top‑p 篩選)回傳 logprob。PipelineRL 需要的是抽樣器使用的已處理分布的 logprob。透過設定 logprobs-mode=processed_logprobs 解決此問題,消除了 rollout logprob 的平均偏移,並使平均 policy ratio 接近 1.0。
執行時預設值
為確保等價,團隊明確停用與 V0 參考路徑不同的 V1 特定預設值:
- 前綴快取:停用 (
enable-prefix-caching: false) 以防止在權重更新前計算的狀態被重複使用,這會在快取壽命上產生僅 V1 的差異。 - 非同步排程:停用 (
async-scheduling: false) 以消除另一項僅 V1 的自由度。
進行中權重更新
為了匹配 V0 在載入新權重時不需明確失效快取狀態的行為,團隊實作了以下 V1 更新路徑:
await engine.pause_generation(mode="keep", clear_cache=False)
await engine_client.collective_rpc_async(
"receive_weight_update",
args=(request.model_dump_json(),),
)
await engine.resume_generation()
使用 mode="keep" 與 clear_cache=False 確保權重同步模型與 V0 參考相符,減少了訓練後期的持續延遲。
數值等價:fp32 lm_head
即使在後端修正之後,最終的等價仍需匹配計算 logits 的數值路徑。訓練器使用 fp32 的 lm_head 進行最終投影,且 rollout 後端必須更新以符合此精度。
此精度至關重要,因為 RL 更新直接使用 token logprob;logits 的微小變化會顯著影響 policy ratio 與 KL 散度。此發現與 MiniMax-M1 技術報告以及 ScaleRL 論文相呼應,兩者皆指出 fp32 head 計算是大規模 RL 防止訓練/推論 token 機率不匹配的必要設計選擇。
為何後端正確性優先於目標修正
雖然截斷重要性抽樣或重要性比重重新加權等工具可校正陳舊或非同步的 rollout,但在修正後端之前使用它們會混淆結果。若在目標端套用修正以彌補破損的推論‑後端行為,訓練曲線將變得更難解讀。
團隊的做法確保任何剩餘的不匹配皆源於 RL 目標的設計(例如非同步/離策略問題),而非推論引擎的錯誤。未來的改進將聚焦於非同步/離策略的清理,例如保留 rollout 時的明確 behavior‑policy logprob,並在最佳化時重新計算訓練端的 old‑policy logprob。