返回首頁

訓練基礎設施

TRL 1.14 只同步 LoRA 權重,讓非同步 GRPO 跨機訓練不再依賴 NCCL

Hugging Face 將 `AsyncGRPOTrainer` 的訓練與 vLLM rollout 分散到不同機器,只透過共用物件儲存傳遞數 MB 的 LoRA adapter。公開實驗把 500 步流程由 3 小時 27 分縮至 53 分,但結果仍是單一小模型配方,不等於已證明大規模 RL 收斂或成本優勢。

Mark Harkin · CC BY 2.0 · Image source
zh-Hant

Hugging Face 在 TRL 1.14 為實驗性的 `AsyncGRPOTrainer` 加入 adapter-only 同步:訓練器以 FSDP2 更新 LoRA,隔數個 optimizer step 將新版本寫入 Storage Bucket,再要求獨立的 vLLM 工作透過 `/v1/load_lora_adapter` 載入。以 Qwen2.5-Math-1.5B、rank-1 LoRA 為例,每次搬動的是數 MB adapter,而不是約 3 GB 的完整模型,因此訓練與生成程序不必位於同一節點,也不用建立跨節點 NCCL group。

這套設計的難點不是單純共享檔案。前置 proxy 會補上 HF Jobs 的授權標頭、把同一 rollout 導向較可能保有其 KV prefix 的 replica,並向所有 replica 廣播 adapter 載入、暫停與恢復。每次更新都使用新的 policy 名稱,避免舊權重建立的 KV block 在熱更新後被錯誤重用。若允許樣本落後四個 policy 版本,vLLM 還必須保留目前版本、四個舊版本以及交換期間的額外槽位,因而設定 `max_loras=6`;槽位不足可能在仍有 rollout 執行時逐出舊 policy。

團隊逐步調整 token-budget packing、併發量、gradient checkpointing 與 adapter 載入重試,把同一個 500-step 配方從 3 小時 27 分降至約 53 分,並公開執行腳本及 proxy 測試。這證明低秩更新可把同步問題改寫成版本化成品傳遞,適合沒有高速叢集互連的工作環境;它並非「沒有網路」,而是以 FUSE-backed bucket 與 HTTPS 控制流取代直接 tensor collective。

工程團隊下一步應檢查物件儲存的可見性延遲、adapter 原子發布、policy staleness 與被丟棄 rollout 的比例,也要分開量測 GPU 費用、最終 reward 及收斂穩定性。官方數字僅來自 1.5B 模型和特定 H200 配置,尚未與成熟的 NCCL/Slurm 叢集做等品質、等成本比較;DoRA、額外可訓練模組或超出 vLLM rank 上限的設定也會退回完整合併權重同步。

來源

  1. Async GRPO with LoRA across HF Jobs: a bucket, a proxy, and no NCCL
  2. hfjobs-lora-buckets companion implementation