DeepSeek同款GRPO训练大提速!魔搭开源全流程方案,支持多模态训练、训练加速和评测全链路

魔搭社区 投稿
量子位 | 公众号 QbitAI

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的范式也逐渐引起重视,团队也会在该领域积极探索最新的评测基准、指标和方法。

[1] AI-MO/NuminaMath-TIR数据集:https://www.modelscope.cn/models/AI-ModelScope/NuminaMath-7B-TIR
[2] CLEVR-70k-Counting:https://www.modelscope.cn/datasets/okwinds/clevr_cogen_a_train
[3] 多模态GRPO最佳实践:https://github.com/modelscope/ms-swift/blob/main/docs/source/BestPractices/GRPO%E5%A4%9A%E6%A8%A1%E6%80%81%E8%AE%AD%E7%BB%83.md
[4] 大模型评估框架EvalScope:https://github.com/modelscope/evalscope
[5] modelscope/R1-Distill-Math-Test数据集:https://modelscope.cn/datasets/modelscope/R1-Distill-Math-Test
[6] EvalScope模型推理能力评测最佳实践:https://evalscope.readthedocs.io/zh-cn/latest/best_practice/deepseek_r1_distill.html
[7] DeepSeek-R1-Distill-Qwen-7B :https://modelscope.cn/models/deepseek-ai/DeepSeek-R1-Distill-Qwen-7B
[8] 模型思考效率评测最佳实践:https://evalscope.readthedocs.io/zh-cn/latest/best_practice/think_eval.html
[9] GRPO完整实验流程:https://github.com/modelscope/ms-swift/blob/main/docs/source/BestPractices/GRPO%E5%AE%8C%E6%95%B4%E6%B5%81%E7%A8%8B.md 

—  —


学术投稿请于工作日发邮件到:

ai@qbitai.com

标题注明【投稿】,告诉我们:

你是谁,从哪来,投稿内容

附上论文/项目主页链接,以及联系方式哦

我们会(尽量)及时回复你

一键关注 👇 点亮星标

科技前沿进展每日见

一键三连「点赞」「转发」「小心心」

欢迎在评论区留下你的想法!

(文:量子位)

欢迎分享

发表评论