Hugging Face Jobs 上的 Async GRPO 與 LoRA：無需 NCCL 的分散式訓練新架構

重點摘要

AsyncGRPOTrainer 現在可以訓練 LoRA 適配器，並且只將該適配器同步到 vLLM (TRL v1.14)。一個 rank-1 適配器只有幾 MB，因此可以透過掛載到每個 Job 的 Storage Bucket 傳輸，而非透過 NCCL。訓練器和 vLLM 副本作為獨立的 Hugging Face Jobs 在不同的機器上運行。

副本前的一個小型代理伺服器會添加驗證標頭，將每個推演路由到已持有其 KV 前綴的副本，並將適配器載入廣播到每個副本。AsyncGRPO 的指標顯示了瓶頸所在。五次運行將相同的配方從 3 小時 27 分鐘縮短到 53 分鐘，完成 500 個步驟。

LoRA 支援最近透過 PR #7017 登陸 TRL 的 AsyncGRPOTrainer，並隨 TRL v1.14 版本發布。這個非同步訓練器現在可以訓練一個適配器而非完整的模型，並且只將 LoRA 適配器同步到 vLLM。本文介紹了一個基於此技術的實際專案，其中訓練和推論不再共用同一台機器。

LoRA 訓練特別適合強化學習 (RL)，正如 Thinking Machines 的部落格文章《LoRA Without Regret》所示。他們展示了 LoRA 即使在 rank 1 的情況下，也能在策略梯度 RL 中達到與完整微調相同的效果。

這源於優勢函數 (advantage function) 每個回合只提供約 O(1) 位元的資訊，因此從總資訊量的角度來看，每一步可學習的內容並不多。一個 rank-1 適配器有足夠的容量來吸收這些資訊。

LoRA 訓練也帶來了系統層面的影響。一個 1.5B 模型 的 rank-1 適配器只有幾 MB，而完整模型約為 3 GB。我們現在可以只傳送適配器，而非在每次更新後將完整的策略傳送給推論工作者。vLLM 也能同時載入多個適配器。舊的推演會以其開始時的策略完成，而新的推演則使用最新的策略。

TRL 的 AsyncGRPOTrainer 已經將訓練和生成分開。訓練器和 vLLM 可以在不同的機器上以各自的速度運行。這在單節點或叢集環境中很容易實現，因為兩個程序可以共用檔案系統或組成 NCCL 群組。

我們希望在 Hugging Face Jobs 上運行相同的設定。本質上，一個 HF Job 就是一個容器運行在一個虛擬機器 (VM) 上。這意味著一個 Job 無法產生多個節點 (至少目前如此) 來容納一個訓練器和一組 vLLM 伺服器 (每個節點最多限制為 8 個 H200)。

AsyncGRPOTrainer 正是為這種規模而設計的，所以問題變成了：如果我們放棄訓練器和推論伺服器共用一個節點的要求，我們能走多遠？如果是完整權重同步，答案將是「走不遠」。每次更新都必須在機器之間移動數 GB 的資料，這正是 NCCL 在密集叢集中的作用，但 Jobs 無法跨節點通訊。沒有共用的本地磁碟，顯然也沒有共用的 localhost。

使用 LoRA，同步只需幾 MB。對於檔案系統部分，HF Jobs 提供了由 Storage Buckets 支援的儲存卷！這些儲存桶可以作為 FUSE 檔案系統掛載到每個 Job 中，足以作為節點之間的共用檔案系統。Jobs 之間完全不需要網路路徑。

最終的設定非常精簡：一個運行 AsyncGRPOTrainer 與 LoRA (以及 FSDP，稍後詳述) 的訓練器 Job；兩個 vLLM Jobs，每個都提供基礎模型以及訓練器最新發布的適配器；一個 Storage Bucket 以相同路徑掛載到所有三個 Jobs 中，這是適配器從訓練器傳輸到伺服器的方式；以及一個代理伺服器。

我們將深入探討為何需要它，但從高層次來看，我們需要一個代理伺服器，它將每個推演路由到最有可能持有其 KV 快取的副本，並將每次適配器更新廣播到所有 vLLM 副本。

架構：利用 Hugging Face Jobs 和 Storage Buckets 🪣

AsyncGRPOTrainer 中新的僅適配器同步路徑運作方式如下。訓練器不向 vLLM 傳送張量。每隔幾個優化器步驟，它會將適配器儲存在 <output_dir>/.vllm_lora/trl-policy-v{N} 下，透過原子重新命名發布該目錄，然後將其路徑傳送給 vLLM 的 /v1/load_lora_adapter 端點。

vLLM 從磁碟載入檔案，因此推演工作者可以請求 model="trl-policy-v{N}"。這就是 vLLM 中執行時適配器載入的運作方式。該端點接受路徑而非張量，因此訓練器和伺服器預期會共用一個檔案系統。在 Slurm 叢集中，這是網路檔案系統。

在 Jobs 上，我們透過將 Storage Bucket 作為儲存卷以相同路徑掛載到每個 Job 中來實現相同目的，如前所述。在底層，它使用 hf-mount，將儲存桶作為 POSIX 檔案系統暴露在容器內部：

# 每個 Job 都會以相同的絕對路徑取得相同的儲存桶

hf jobs run ... -v hf://buckets/aminediroHF/asyncgrpo-lora-buckets:/lora ...

TRL 或 vLLM 中沒有任何東西需要為此改變。訓練器寫入 /lora/<run>/.vllm_lora/，伺服器從相同路徑讀取。POST 請求中傳送的路徑在每個容器內部都已有效。

三個 Job 和儲存桶。TRL 透過 localhost 與代理伺服器通訊，代理伺服器透過 HTTPS 與副本通訊，適配器目錄則透過儲存桶掛載傳輸。請注意，我們也將檢查點和最終適配器儲存在儲存桶中。HF Jobs 是短暫的，但被搶佔的訓練器可以恢復訓練，因為最終適配器始終會持久化到儲存桶中，並且在 Job 停止時不會丟失。

三個 Jobs

vLLM 副本

每個副本使用一個 GPU 和標準的 vllm/vllm-openai 映像。我們只需要啟用執行時 LoRA 載入並保留足夠的適配器槽位。適配器槽位的數量取決於 max_staleness。在 AsyncGRPOTrainer 中，每次權重同步都會將策略版本增加一，而 max_staleness 是推演樣本可能落後於當前策略多少個版本，訓練器才會將其丟棄。

當訓練器處於 v7 時，在 trl-policy-v3 下生成的樣本仍用於訓練，如果 max_staleness=4。在 v3 下開始的推演也必須能夠在 v3 下完成。因此，在任何時刻，vLLM 都必須提供當前策略以及之前的四個策略。這就是為什麼訓練器會註冊 max_staleness + 1 個適配器版本，並卸載任何更舊的版本。

每次同步都會在卸載最舊版本之前載入新版本，這在交換期間需要額外一個槽位。這就是為什麼需要 --max-loras 6。如果只有五個，vLLM 會在每次同步時默默地驅逐一個仍有推演正在進行的策略。

# --expose 8000 可透過 https://<job_id>--8000.hf.jobs 存取

# -v ...:/lora:ro 唯讀：伺服器只讀取適配器

# VLLM_ALLOW_RUNTIME_LORA_UPDATING=1 啟用 /v1/load_lora_adapter

# VLLM_SERVER_DEV_MODE=1 啟用 /pause, /resume, /server_info (TRL 需要這三個)

# --max-loras 6 max_staleness=4 -> 4+2 適配器槽位

for replica in 1 2; do

hf jobs run --detach --flavor h200 --timeout 8h --secrets HF_TOKEN \

--expose 8000 \

-v "hf://buckets/${BUCKET}:/lora:ro" \

-e VLLM_ALLOW_RUNTIME_LORA_UPDATING=1 \

-e VLLM_SERVER_DEV_MODE=1 \

-- vllm/vllm-openai:v0.27.1 \

vllm serve Qwen/Qwen2.5-Math-1.5B --host 0.0.0.0 --port 8000 \

--max-model-len 4096 --logprobs-mode processed_logprobs --generation-config vllm \

--enable-lora --max-lora-rank 1 --max-loras 6

done

我們將 vLLM 固定在 v0.27.1 版本。vLLM 發展迅速，上述旗標和執行時 LoRA 端點是該版本所提供的，因此請將此版本視為配方的一部分。還有另一種可能的設計，即訓練器只保留最新的適配器並始終以相同的名稱發布。我們沒有選擇這種方式，因為 vLLM 以適配器名稱作為其前綴快取的鍵。

使用單一名稱，在交換後，在先前權重下計算的 KV 區塊仍會匹配，因此預填充 (prefill) 不會重新執行，推演可能會從一個策略版本取得其前綴，並從下一個版本取得其解碼。訓練器將無法判斷，這將表現為比例偏離 1。版本化的名稱使這成為不可能：一個名稱始終代表一組權重，並且快取的前綴永遠無法匹配更新的版本。

資料集選擇：Sanity 集

我們選擇了 sail/Sanity-Test-R1D-1.5B，這是來自《Defeating the Training-Inference Mismatch via FP16 (Qi et al., 2025)》論文的資料集。重現程式碼位於 sail-sg/Precision-RL。

作者使用 DeepSeek-R1-Distill-Qwen-1.5B 為每個數學問題生成了 40 個答案。他們保留了成功率在 20% 到 80% 之間的問題，產生了 1,460 個問題。

這個資料集非常適合 RL 驗證，因為這些問題對於該模型來說既非已解決也非完全無望，這意味著模型可以獲得良好的早期訊號進行訓練和改進。這作為一個穩健的端到端測試非常棒：如果一個 vLLM 副本在適配器名稱下默默地提供基礎模型，我們希望在幾十個步驟內在曲線上看到這一點。此外，這個資料集足夠小，可以在不到兩小時內完成循環。

我們還採用了論文中 oat/scripts/lora 裡的 LoRA 腳本的超參數：Qwen/Qwen2.5-Math-1.5B，LoRA rank 1，alpha 2，學習率 4e-5，每個提示詞 8 個樣本，每步驟 128 個完成，最大生成 3,000 個 token，以及 4,096 個 token 的上下文。

訓練器

訓練器使用相同的 vllm/vllm-openai:v0.27.1 映像，並在其上安裝 TRL。我們當時運行的是 PR 分支；相同的程式碼現在隨 TRL v1.14 發布。訓練腳本是一個正常的 AsyncGRPOTrainer 腳本。唯一與 Job 相關的值是輸出目錄和伺服器 URL。

from peft import LoraConfig

from trl.experimental.async_grpo import AsyncGRPOConfig, AsyncGRPOTrainer

config = AsyncGRPOConfig(

output_dir="/lora/sanity-lora-r1", # 在儲存桶上：適配器、檢查點和最終適配器都將儲存於此

vllm_server_base_url="http://localhost:8000", # 代理伺服器，而非 vLLM Job；TRL 從未見過 Jobs URL

max_staleness=4,

weight_sync_steps=4, # 每 4 個優化器步驟發布一個適配器

save_strategy="steps", save_steps=50, # 檢查點儲存到相同的儲存桶 -> 被搶佔後可恢復訓練

...

)

trainer = AsyncGRPOTrainer(

model="Qwen/Qwen2.5-Math-1.5B",

args=config,

peft_config=LoraConfig(r=1, lora_alpha=2, target_modules="all-linear"), # vLLM 可直接提供純 LoRA

...

)

在初始化期間，TRL 會呼叫 /server_info。如果它找到 lora_config，就會使用僅適配器同步。vLLM 無法直接提供的配置，例如 DoRA、modules_to_save 或超過 --max-lora-rank 的 rank，將會回退到合併權重同步並發出警告。日誌中應包含 "Adapter-only vLLM sync enabled"。

代理伺服器

現在來談談有趣的部分。我們需要在訓練器和 vLLM Jobs 之間設置一個代理伺服器，原因有二：暴露的 Job 埠需要每個請求都帶有 Authorization: Bearer <HF token> 標頭。代理伺服器就是添加這個標頭的地方，因此 TRL 不需要知道它。

我們希望有多個 GPU 進行生成。在單一 vLLM 伺服器上，通常的做法是使用 --data-parallel-size > 1，但 TRL 在此模式下拒絕僅適配器同步，原因很充分：呼叫 /v1/load_lora_adapter 只會到達響應它的 DP rank，因此其他 rank 將繼續在新策略名稱下提供基礎模型。

在 Jobs 上，這個問題甚至不會出現，因為每個副本都是獨立的機器。因此，資料並行性必須提升一個層次，由某個東西將適配器載入分散到每個副本。

因此，我們在訓練器 Job 上運行一個小型代理伺服器在 127.0.0.1:8000，並將 TRL 指向它，就好像它是一個單一的 vLLM 伺服器一樣。除了添加標頭之外，代理伺服器在功能上還做兩件事：它將每個完成請求發送到一個副本，選擇的原則是讓一個提示詞的八個推演落在其前綴已經快取的地方 (詳情如下)。

它將每個改變狀態的請求，例如適配器載入、暫停和恢復，廣播到所有副本，以便策略名稱在所有副本中都意味著相同的事情。