DeepSeek同款GRPO訓練大提速!魔搭開源全流程方案,支持多模態訓練、訓練加速和評測全鏈路
GRPO訓練又有新的工具鏈可以用,這次來自于ModelScope魔搭社區。
隨著DeepSeek-R1的成功出圈,其使用的GRPO算法受到了業界的廣泛關注。GRPO訓練是來自于PPO算法的一種改進,旨在利用采樣原理對value model進行簡化,以增大訓練的穩定性和可維護性。
目前圍繞R1模型的技術方案,社區也已經有一些開源實現,比如來自Hugging Face的Open-R1,以及其他一些RL框架,包括veRL,OpenRLHF等等。然而,在GRPO訓練方面,大部分方案仍然面臨著諸多挑戰,包括訓練速度較低、集群配置復雜,以及對多模態擴展能力差、訓練后不好評測等等。
為了支持開源社區在GRPO這個方向上的探索,ModelScope魔搭社區圍繞MS-SWIFT訓練框架以及EvalScope評估框架,推出了相對完整高效的GRPO全鏈路解決方案,和社區分享。
GRPO訓練提速
GRPO訓練的耗時主要來自于采樣、訓練等幾個方面。其中,采樣本身是PPO算法的重要組成部分。尤其是GRPO在PPO的基礎上,使用采樣代替value model的方案,這使得在GRPO訓練中,采樣的耗時的占比,更是大大的增加了。而且GRPO中單query的采樣數(即group size),一般比較大(DeepSeekMath論文中為64個),高采樣頻率對于推理引擎的挑戰是巨大的。優化采樣效率是提升GRPO整體訓練速度的核心所在。
基于這些特點,魔搭社區的SWIFT框架進行了有針對性的優化:
多實例數據并行采樣
對于GRPO算法,單實例采樣往往無法滿足需求。團隊觀察發現,7B模型在訓練期間,單iter的采樣時間占比約為70%,這意味著應該允許訓練資源根據實際情況,針對性的傾斜分配給采樣計算。
尤其在采樣量以及batch_size較大時,采樣耗時對訓練速度影響會更為顯著。因此,在SWIFT中對vLLM和LMDeploy進行了一定的patch(同時也與相關框架負責同學進行了討論,將會將相關實現在vLLM/LMDeploy上原生支持),支持在任意比例的訓練卡上拉起采樣實例。例如對于8卡訓練中,配置4張卡負責模型訓練,4張卡負責采樣;或者6張卡負責訓練,2張卡負責采樣。
下圖展示了在同樣的8卡訓練設置下,分別使用1卡/2卡部署推理引擎進行采樣,剩余卡負責訓練,vLLM/LMDeploy的采樣耗時和訓練耗時如下:
可以看到,基于LMDeploy分配2卡進行采樣時,訓練時長約為1卡采樣的87%。而vLLM2卡采樣時,時長為1卡采樣的78%。在這兩個例子中,通過更合理多卡數分配,對采樣資源進行一定傾斜,都達到了有效降低訓練所需時長都目的。
異步采樣
GRPO訓練中,采樣和模型訓練需要交替進行,即訓練時采樣卡閑置,采樣時訓練卡閑置。對于這個問題,不同的框架給出了不同的解決方案。
例如veRL允許將采樣和訓練部署在同一GPU上,在不同階段讓不同的權重offload到CPU中;或者針對LLM的不同Layer、不同Tensor進行異構式切分,在加載權重時不必all_gather(而是部分weights進行gather并同步),使訓練模型和采樣模型的效率達到最高。然而在中小模型的體量上,這樣的切分未必是最優的。
因為隨著模型尺寸和batch_size增大,采樣、訓練的耗時占比會有本質差別。對此SWIFT采用了不一樣的技術路線,即異步采樣(replay buffer),其原理是在訓練時同時進行采樣,采樣結果用于下一iter的模型訓練。由于采樣使用的是old policy model,因此訓練過程需要額外增加對于logits差異的CLIP。由于old policy model和policy model僅相差一個iter,因此其訓練穩定性幾乎沒有下降。二者唯一需要等待(或者stop the world)的過程是權重加載。
同樣訓練設置下的實驗測試表明,在單卡部署LMDeploy的情況下,異步采樣的訓練時間約為同步采樣的2/3。
模型placement
SWIFT除了支持上述訓練和rollout使用兩個資源組進行異步訓采流程之外,也支持二者共用同一資源組。即,在actor模型訓練時,vLLM將開啟sleep模式以減少顯存占用。
這兩種模式的架構圖如下:
并且,SWIFT也支持vLLM的tensor_parallel(MP)模式。
LMDeploy推理框架支持
LMDeploy是來自于上海浦江實驗室的優秀推理加速框架。該框架不僅支持了純文本和多模態模型的推理加速,并且自研了基于FasterTransformer實現的Turbomind推理加速引擎。在推理速度上,LMDeploy相較vLLM在一眾模型上速度有顯著的提升。對Qwen2.5-7B-Instruct模型進行實驗測試,具體配置如下:使用AI-MO/NuminaMath-TIR數據集,批量大小設置為7,每條query采樣24條結果,訓練50個steps。以下是vLLM框架和LMDeploy框架在相同條件下的推理時長對比。
可以看到使用LMDeploy作為采樣的推理引擎,在整體的訓練速度上從44分/50steps加速到37分/50steps,加速比大約為16%。
注:最后50steps的訓練耗時包括了保存模型權重,測試集評測
SWIFT框架在TRL框架以及vLLM采樣的基礎上,額外引入了LMDeploy采樣的支持。得益于其推理速度和權重加載速度,可以做到整體采樣耗時只有基礎實現的70%。
多輪更新
多輪更新的核心思想是:讓一次采樣得到的數據可以被多次使用。通過這種方式,可以減少采樣的頻率,從而讓采樣和訓練的資源分配更加均衡。
通過設置參數num_iterations, 可設置每輪采樣數據的迭代更新輪數,增大該參數值可以多次利用采樣數據,緩解采樣過程對訓練的速度影響,從而提高訓練速度。在該值設置不太大(例如小于等于4)的情況下,基本不會對模型訓練效果有不良影響。這里的更新輪次對應論文中的mu值。
多輪更新是TRL庫提供的機制,發現該機制和團隊提供的其他機制結合起來有更好的加速效果。
基于1卡部署LMDeploy的設置下,對比num_iterations從1-4的訓練時間,下面給出了完整的實驗結果:
可以看到,在設置多輪更新輪次為4時(mu=4),整體訓練耗時約為單輪更新的一半。
綜合測試
針對八卡環境的SWIFT、veRL和trl框架的訓練效率進行了對比。在實驗設置中整合了前文提到的多種訓練加速技術,并選用LMDeploy作為推理引擎。
具體配置上,采用雙卡進行推理采樣,結合異步采樣策略,并將多輪更新輪次設定為4輪。同時,為了更好模擬真實訓練場景,將batch_size調整為48(每輪6個query,梯度累計步數為 8),group size設置為24,基于Qwen2.5-7B-Instruct模型和AI-MO/NuminaMath-TIR數據集(1),對多框架的GRPO訓練速度進行了對比評測。
SWIFT:
veRL:
trl(mu=4):
trl(mu=1)
實驗結果表明,SWIFT框架的訓練耗時約為120秒/步,而veRL框架訓練耗時約為280秒/步,TRL框架在采用多步更新下約為144秒/步,而不采用多步更新則約為320秒/步。通過整合多種訓練加速技術,SWIFT框架在GRPO的中小集群訓練效率上實現了顯著提升。下圖展示了SWIFT框架下的訓練獎勵走勢,可以看出模型成功實現了獎勵值的提升。
多模態GRPO訓練
多模態GRPO訓練目前已經有一些開源解決方案,例如R1-V,open-r1-multimodal等,都是基于Open-R1進行了簡單的擴展。
SWIFT框架目前已經支持多模態模型(圖文、視頻、音頻)的GRPO訓練,只要在數據集中給定’images’/‘videos’/‘audios’字段,GRPO會將多模態內容輸入多模態模型中進行強化訓練。SWIFT目前在微調中支持的多模態模型近兩百個,這些模型均天然支持GRPO訓練。參考R1-V的任務設置,在多模態計數任務上進行了訓練,訓練數據集為CLEVR-70k-Counting(2)。訓練選取的獎勵函數一共有兩個,一個是 Deepseek-R1 中提到的格式獎勵函數,用于評估模型輸出格式的準確性;另一是自定義的準確性獎勵函數,用于計算模型輸出的計數是否與真值一致,目前兩個獎勵函數都已經定義在SWIFT框架中,通過—reward_funcs external_r1v_acc format參數指定。
選取Qwen2.5-VL-3B-Instruct 作為基礎模型進行訓練,選取 instruct 而不是base模型的主要原因是可以更快地獲取 format reward。整體實驗在8卡 GPU 上完成。當前SWIFT GRPO訓練已支持多卡部署模型以加速rollout,因此設置num_infer_workers為2,進程數為6:即2卡vLLM部署采樣,6卡模型訓練。設置模型的最大輸出為1024,學習率設置為1e-6,其他參數設置詳見最佳實踐(3)。
實驗結果如下圖所示:
模型訓練了500個epoch已經基本收斂,準確性獎勵(圖中為ClevrCountORM)和格式獎勵(圖中為Format)在不斷增加,證明模型學習到了如何完成這一任務,最終任務成功率從初始的0.4攀升到1左右;在大概300step左右,reward_std已經將到0.1左右,證明模型已經基本收斂;completion length最終穩定在60-80,模型學習到的任務推理范式是逐個列舉圖中的物體。
訓練后的模型輸出樣例如下:
user:
How many items are there in the image?
assistant:
<think> Counting the number of items in the image:\n1. Green matte sphere\n2. Large metallic yellow sphere\n3. Small metallic brown cube\n\nThere are three distinct objects in total.\n</think>\n<answer>3</answer>
Reasoning模型評測
EvalScope框架是魔搭社區上開源的大模型評測工具(4),提供了完整的大模型全面評測框架。
不僅提供O1/R1等推理類模型的推理性能的評測能力,還支持評測結果的可視化,如下圖所示。
同時,團隊將MATH-500、GPQA-Diamond和AIME-2024三個數據集整合為一個數據集合,放置于modelscope/R1-Distill-Math-Test數據集(5)中,用戶可以直接使用該數據集的ID進行評測操作,具體使用流程參考:模型推理能力評測最佳實踐(6)。
除此之外,針對Reasoning模型在推理過程中存在Underthinking(思考不足,指模型在推理時頻繁思路跳轉,無法專注于正確的思路,導致錯誤答案)和Overthinking(過度思考,指模型簡單問題上生成過長思維鏈,浪費計算資源)的問題。
該框架實現了模型思考效率評測能力,可以評估如DeepSeek-R1-Distill-Qwen-7B(7)等推理模型的思考效率,如下圖所示,從token效率、思考長度、子思維鏈數量和準確率四個方面進行衡量,可以用來評測并優化模型在短輸出中獲取正確答案的能力,具體使用流程參考教程:模型思考效率評測最佳實踐(8)。
效果
從簡單的數學任務Countdown Game出發,對SWIFT框架的GRPO有效性進行驗證,并給出了完整的實驗流程(9)。
Coundown Game 的任務目標是根據給定的幾個數字和加減乘除四種運算,得到目標數字,給出運算公式,因此在模型的輸入中給定了任務描述、已有數字和目標數字。訓練選取的獎勵函數一共有兩個,一個是標準的格式獎勵函數,用于評估模型輸出格式的準確性;另一是自定義的準確性獎勵函數,用于評估模型輸出運算公式是否能得到目標值,目前兩個獎勵函數都已經定義在SWIFT框架中,通過—reward_funcs external_countdown format參數指定。
選取Qwen2.5-3B-Instruct模型進行全參數GRPO訓練,具體參數設置詳見實驗流程。發現GRPO的訓練過程是比較不穩定的,一個比較大的學習率和beta(KL散度的權重)雖然會是模型更快收斂,但往往出現梯度過大而訓練崩潰的情況。最終選取學習率為7e-5,beta為0.001。
實驗結果如下圖所示:
一共訓練了2000個step,各項曲線證明了SWIFT GRPO的有效性,從實驗結果中也能看到模型的收斂過程。
準確性獎勵(CoundownORM)和格式獎勵(Format)的穩步上升證明模型逐步學習到了如何更好地完成這一任務,同時reward_std最終處在0.2-0.3左右也證明了模型仍有繼續上升的空間。
從completion_length的變化上能明確地看到模型思考方式的轉變。從初始的500左右降到200,然后又增長到300-400 tokens,能明確看到各個階段模型的思考方式。初始時,模型會直接從結果出發,反推需要的數字;然后,模型逐漸精簡了反推過程,顯著縮減了輸出場地;最后,模型發現完成任務最好的方式是列舉已有數字和運算的組合,并逐步運算,輸出范式基本收斂。最終的輸出樣例如下:
user:
"Using the numbers [71, 92, 80], create an equation that equals 59. You can use basic arithmetic operations (+, -, *, /) and each number can only be used once. Show your work in <think> </think> tags. And return the final equation and answer in <answer> </answer> tags, for example <answer> (1 + 2) / 3 = 1 </answer>."
assistant:
"<think>\nWe need to form an expression using the numbers 71, 92, and 80 which results in 59. Let's try some combinations:\n(92 - 80) + 71 = 83 + 71 = 154 (too high)\n(92 - 71) + 80 = 21 + 80 = 101 (too high)\n(92 - 80) + 71 = 12 + 71 = 83 (too high)\n(92 - 71) + 80 = 21 + 80 = 101 (too high)\n(92 + 80) - 71 = 172 - 71 = 101 (too high)\n(92 + 71) - 80 = 163 - 80 = 83 (too high)\n(80 + 71) - 92 = 151 - 92 = 59\n\nSo our solution is: (80 + 71) - 92 = 59</think>\n\n<answer>(80 + 71) - 92 = 59</answer>"
寫在最后
SWIFT正在添加vLLM的MP模式以適配更大體量模型的訓練,并希望在已有的優秀框架的技術沉淀上以差異化技術的方式為開發者提供簡單快速的中、小集群RL訓練方案,給開發者提供新的技術選型。目前SWIFT已經支持的訓練領域有數學、ReACT結構的Agent、多模態VQA等,代碼領域的支持也在陸續更新中。對于Megatron結構的模型SWIFT在支持中,不僅僅是GRPO訓練,SFT和PreTrain也將得到相應支持。
評測領域,EvalScope后續會針對Reasoning models進一步探索模型的“思考效率”問題;另外,從目前的趨勢看,多模態Reasoning的范式也逐漸引起重視,團隊也會在該領域積極探索最新的評測基準、指標和方法。