在人工智能技术的快速发展中，多模态大语言模型（MLLM）以其强大的图文理解、创作、知识推理及指令遵循能力，成为了推动数字化转型的重要力量。然而，如何使这些模型的输出更加贴近人类的风格、符合人类的偏好，甚至与人类价值观保持一致，成为了一个亟待解决的问题。为了应对这一挑战，基于人类反馈信号的强化学习方法（RLHF）应运而生，其中，PPO（Proximal Policy Optimization）算法作为 OpenAI 的核心技术，在 RLHF 阶段扮演着关键角色。
小红书大模型团队，在这个技术日新月异的时代，开始了他们自研 MLLM RLHF 训练框架的征程。他们深知，要构建一个高效、准确的 RLHF 训练系统，需要综合考虑算法优化、系统架构、训练调度以及推理引擎等多个方面。
1.1 RL强化学习原理
在介绍本文工作之前，我们先了解一下RL强化学习的原理。
状态S：输入prompt
动作A：输出response
奖励R：根据prompt+response进行奖励模型打分
给定prompt，调整policy，生成符合人类喜好（RM偏序信号）的response
RL强化学习用在游戏、自动驾驶等场景较多，简单解释就是智能体在特定环境下进行不同策略的action、并给出反馈进行状态更新的过程，可以理解为自闭环；对应到NLP下，大语言模型就是智能体，输入的提示词prompt就对应状态S，输出下一个token就是动作A，对prompt+res进行打分就是奖励R，整体目标就是给定prompt，调整policy策略模型，生成符合人类喜好的response，这个喜好就是RM(Reward-model)的偏序信号。
1.2 RLHF-PPO算法
RLHF的流程主要是RM+RL两个阶段，sft是RL的前置阶段，一般作为RM/RL的初始化模型。
第一阶段是RM，偏好训练一个reward-model打分模型，这里的数据是有偏的，目的就是模拟人类视角去评价一条数据的质量。
第二阶段是RL，LLMs根据训练集的指令生成自己的偏好回复，并根据RM的实时打分，进行策略优化，不断生成接近人类的偏好输出。
RL算法一般有on-policy、off-policy两类，OpenAI在RLHF阶段采用的有效算法是PPO，PPO是一种on-policy算法，需要在线生成样本，被普遍认为是效果较好的一种RL算法，所以我们RL训练也是基于PPO进行设计和优化的；
RL训练-PPO（Proximal Policy Optimization策略梯度算法）：
1.Actor Model：演员模型，训练的目标语言模型
2.Critic Model：评论家模型，预估总收益
3.Reward Model：奖励模型，计算即时收益
4.Reference Model：参考模型，“约束”，防止模型训歪
PPO是一种经典的策略梯度算法，具体算法涉及到actor/critic/reward/reference 4个模型的协同训练和推理，设计和实现一套高效、准确的RLHF训练系统，是关键挑战之一。
PPO训练流程
下面介绍一下PPO的训练流程，PPO是一个训推混合的训练流程，如下图所示，可以分为两个步骤，一个是make-experience，一个是training。
第一阶段经验采样，有时候也叫做rollout，就是滑跑的意思，在RL中可以理解为一次实验，也就是actor模型在当前环境下进行策略动作的过程，具体实现上就是actor模型根据prompt数据集进行generate生成response，然后根据prompt+response进行forward计算，得到logp/values/reward等元素，这里涉及到actor、ref/critic/reward 4个模型的推理过程；
第二阶段就是训练流程，涉及到actor和critic两个模型，从计算loss来看，算法上是相互独立的，本质上是两个模型独立训练；
工程实现上的痛点也比较明显：涉及到多阶段的dataloader（prompt和rollout、train），actor模型的generate自回归过程，以及4个不同的模型推理，2个不同的模型训练。
当然整体来看make-experience阶段部分就给我们提供了一定的优化空间。
2.1 整体架构
架构选型
训练：
apex/nemo/megatron-lm/pytorch-lighting/accelerate/ray
=> Megatron-core+ray
推理：
vLLM推理引擎（logp用megatron）
下面介绍我们的整体系统架构，框架选型上我们吸取了开源比如trlx/openrlhf的痛点，比如apex/nemo/megatron等混合框架的组合缝合怪方式，这会给代码维护和更新增加成本，因此我们降维到megatron-core作训练，ray做调度的方式降低维护成本；推理服务我们使用vllm推理引擎，当然需要注意的是，vllm也可以输出训练所需要的logp数据，但还是要和训练保持一致，走megatron框架，否则会引入框架之间的diff，成为训练的bias，引入风险；
该图为我们的整体架构，上面是make-experience的第一步骤，dataloader迭代器返回prompt之后，通过vllm的generate输出response，prompt和response组合之后成为新的forward数据集；左下方是第二步骤，也就是进行4个模型的forward，生成最终的ppo-elements或者简称rollout，构造最终的train数据集；右方为PPO训练，actor训练之后会进行参数同步，这里会涉及到medusa-sft的协同训练流程，最终同步给vLLM进行新一轮的generate等。
2.2 异构组网架构
从0到1优化设计
1.Forward-Offload：纯串行实现，分时复用 model：4->2
2.Training-Seperated：actor/critic 独立集群异步训练（并行化）
启动流程：actor集群上进行torchrun启动主进程训练，并自动以mock方式拉起critic所在集群——remote-func
下面介绍异构组网架构，RL的workload涉及到4个模型，如果全都串行实例化每个model，成本较高，但根据前述分析来看模型结构相同，所以我们设计了异构组网架构，能够大幅提升性能；
如上图所示：
Forward阶段涉及到的4个模型通过offload的方式可以进行串行实现（actor->offload->ref，critic->offload->RM），能够分时复用集群，整体model Memory从4降到2，以llama3 70B为例，需要8机64卡跑一个模型，那么原始的纯并行方式需要4倍，也就是32机256卡跑一组实验，通过offload的方式可以降低到16机128卡，降低模型所需要的GPU集群规模；
Training阶段实现seperated分离的方式，一方面actor/critic模型需要独立参数，另一方面通过独立集群部署可以实现异步训练，所以通过异步的方式可以保证吞吐最大化，这一版本架构已经具备一定可用性，相比开源框架比如trlx/openrlhf提升50%以上；
启动方式：异步的启动方式是复用torchrun拉起主任务，我们实现mock-run的方式，继承同等规模的actor集群信息，能够在远端自动拉起critic所在的集群。
2.3 同构组网架构
新的瓶颈出现，见招拆招
异构组网的痛点：
数据量，seq，模型参数量，集群规模压力加大（以llama3 70B为例，SFT模型本身需要4机32卡，32k长文下CP2=>8机64卡
rollout同步：序列长度、多模数据量增加，同步的通信耗时增加
AC同构组网架构：AC复用集群，main-model 2->1 ——tride-off
forward阶段：offload param进行模型切换training阶段：offload中间状态（param/grad/ddp-bucket/opt）
接下来介绍同构组网架构，当数据量、模型规模等都增大的情况下，集群规模压力也大幅增大，我们以llama3 70B为例，SFT模型本身需要4机32卡，32k长文下需要开启CP2，也就是需要8机64卡，即便是DP1，原始异构组网也需要16机128卡，这个对于资源的压力是较大的；另外异构组网因为需要master-worker架构的拉远实现，需要同步rollout的buffer数据，当多模数据增加了图片等数据后，会大幅增加rollout同步所需要的开销；
所以我们进一步实现了异步组网架构，前面讲到forward阶段的2个模型的同构，训练阶段是异步的，但其实actor和critc部分也是类似，可以通过offload完成进一步的scale，集群规模降低到单个sft所需的训练规模，具体实现上，forward阶段只需要offload-param即可完成模型切换，training阶段，需要切换较多的状态，包括param、grad/ddp-bucket/opt-state等所有显存开销,虽然会增加上下文切换的开销，但相比训练本身的时间相对较小；
这里列举了方案对比，AC异构训练，吞吐最高，但所需集群规模较大，AC同构组网通过完全复用集群的方式，集群资源最小，且最大化性能。
3.1 训练性能优化
数据加载：prefetch，双dataloader（ptx-loss）
并行优化：TP/PP/CP/SP
显存优化：recompute，大模型+长文本
Dynamic-batch：推理、训练不同batchsize
负载均衡：prompt请求Round-Robin到vllm-engines
vllm优化：调整切分，虽降低running-batch，能提高并发
首先是一些常规的训练性能优化方法，数据加载部分采用prefetch，比如ptx-loss的时候存在双dataloader，这里就需要降低num-workers数，防止CPU-OOM；并行策略上，使用megatron的TP/PP/CP/SP并行方法；在长文和大模型下采用recompute方式进行显存优化；dynamic-batch部分的训练和推理用不同的batchsize，因为推理不需要梯度优化器等显存占用，所以bs可以更大；vllm部分优化采用的是round-robin的方式把prompt请求分发给多个vllm-engines进行负载均衡，并且推理服务的TP并行切分可以更细，比如TP4到TP1，虽然降低了running-batch，但是能够增大vllm并发数，把整体的吞吐进一步提升。
3.2 Pipeline优化
问题：
actor进行采样阶段，训练无需进行，存在较大的空窗
思路：
•训推集群混布，增大serving并发，降低generate耗时
•类比pp并行，把generate和forward阶段进行流水线并行
整体调度优化：
•全量offload：
•vLLM：kvcache销毁，param-offload
•Forward：actor/critic的param/buffer-offload
•资源最大化：vLLM资源合并到train
•流水线优化：Make-experience阶段最大化
进一步分析发现，actor进行采样阶段，训练无需进行，存在资源浪费，即便是通过训推资源的混合切换，也会存在较大的空窗；因此我们设计了pipeline的流水线。主要思路是首先通过训推集群混布，增大serving并发，降低generate耗时，我们实现了vllm和megatron的全量offload；vllm部分支持kvcache的销毁，param-offload，forward部分涉及到actor/critic的param以及其他buffer的offload切换，最终可以把vllm资源合并到train内部，实现资源最大化；并且通过类比PP并行，把generate和forward阶段进行流水线并行,通过细粒度的mini-batch切分workload，大幅降低make-experience阶段的整体耗时。
我们整体对比一下不同的实现方式：
常规版本：每个节点独立运行各自的任务，推理在推理集群上部署，ac的训练和推理也在各自的集群上部署，但是推理服务和训练性能差异较大，存在一定优化空间；
全量offload：相当于vllm/actor/critic分别进行全集群的资源利用，这种方法实现简单，可以充分提高generate的并发能力，roi较高；但是也引入了明显的offload/restart等空隙；
Pipeline版本：通过流水线机制可以overlap掉数据落盘、restart/offload，以及一部分forward等开销，在多模等复杂场景下，可以进一步overlap vit等上下文开销，在开源模型上训练时性能最优。
3.3 PPO细节处理
针对PPO我们做了一些细节处理：
padding-free优化，对generate/forward/train三阶段进行padding的去除，vllm通过running-batch进行排队，forward和train通过megatron的1F1B之前进行broadcast即可保证micro-batch内部的正常运算，整体提升10%；
Vllm和训练部分的参数同步上，通过提前和ray-vllm组建ncclGroup的方式完成，相比走文件系统的reload方式进一步提升几十秒；
CP并行部分与传统sft不同，需要进行logp以及values/reward的重排，比如logp部分，规避对logits进行gather的方式，先统一得到total的logp，这样可以把通信量从bsh降低到bs。
3.4 多模MLLM优化
MLLM （Multimodal Large Language Model）
痛点：
vlm部分需要图像处理，图文混合，负载不均衡
特点：
1.视觉部分计算量大，但参数量较小
2.PPO场景下图像处理存在冗余
优化：
1.组网优化：llm 的 tp/pp/cp 用作视觉部分的 dp
2.多路复用：vllm/actor/critic复用相同的img_feature
3.Prefetch：frozen visual-model，通过prefetch的机制，融合进make-experience做Pipeline的overlap
在大规模的场景下，图文混合场景增大的复杂度，也带来了新的瓶颈，因此需要进行特定的性能优化：
我们观察到视觉部分计算量大，llm部分模型权重大。针对这个现象，可以对视觉部分和llm使用不同的并行策略：视觉部分只有 dp; llm 部分有mp + pp + cp。把llm 部分的pp、mp、cp用作视觉部分的dp；
MLLM针对特定的训练任务进行优化，目前sft/RL阶段均采用frozenvisual-model的方式进行，只是微调adapter部分，所以可以进行多路服用，vllm/actor/critic三者复用相同的img_feature；
并且通过prefetch的机制，类似把vit视觉部分的整个实现放到dataloader内部，通过与make-experience阶段结合Pipeline做overlap，能进一步降低预处理的耗时。
3.5 训推一致性
确定性训练和推理
问题：
1.RM的精度要求——RM-serving，acc掉点
2.推理评估的精度要求
原因：
rm的vhead是一个linear，y=xA^T+b
本质上是一个矩阵乘，[1,hidden] * [hidden,1]
如果logits存在框架等diff，则会放大误差：初始diff ~= 0.001，矩阵乘之后达到0.3
训推一致性是对工程架构的一个考验，比如训练之后的模型走推理框架评估可能会存在打分差异，特别的，算法上严格要求RM（Reward-Model）的精度一致，最早通过RM-serving进行accuracy统计发现掉点，经过排查发现是因为框架的diff，逐位误差在千分位，但是RM的vhead是一个linear，本质上是一个矩阵乘，相当于扩大了hidden平方倍，精度无法容忍；
我们的解决方法如下：
RM+RL的训练均使用相同的训练框架；
RL场景下要把RM-serving走api的方式转为本地offload去做forward，这样能够保证训推完全一致，且没有额外的serving占用和通信开销；
训推一致性，训练和推理采样mcore的计算workload，保证评测任务上训推tokenwise一致，如下图所示，多个评测case上完全相同；这里也会涉及到一些技术实现，比如通过allgather+local-reduce实现替换allreduce，cublas的调优选项对齐等细节处理。
3.6 Medusa提升采样效率
投机采样无损优化
问题：
•actor进行generate采样耗时较为明显
•传统量化等方式有一定的精度损失，要求无损优化
推理优化：
•投机采样能无损优化自回归性能
•计算换空间，FLOPS -> VRAM带宽
•Medusa：多个head，frozen主干
•Running-batch：吞吐量最大化
Rollout阶段的generate耗时始终是一个大头，所以我们结合推理部分进一步优化，但传统量化等方式有一定的精度损失，为了稳定PPO训练，也要要求无损优化。
投机采样是一种无损优化，通过多个头的预测来降低decode次数，因为自回归部分主要为 Memory-bound，进行计算并行优化，也是一种计算换空间的方式。这里我们采用的是medusa算法，由主干网络加上多个medusa-head组成，每个MedusaHead是和LM-Head一样结构，分别对应到预测每个头的topk的token，并且RL场景下需要组batch关注的是吞吐，不需要考虑延迟，所以running-batch尽可能大才能提高整体吞吐量；
伴生训练设计
问题：
Backbone vs mhead的匹配，接收率下降
伴生训练：
1.精度无损：backbone实时更新， mhead update
2.训练overlap：mhead的更新可以延后于主干
3.少量微调：数据量rollout-batch-size
性能说明：
由于RL训练阶段均为组batch，所以提升效果有折损，generate部分性能提升50%+
但RL场景下存在一个问题，actor模型在不断更新，medusa的接收率（猜测准确越高，降低decode次数越多，越能提升推理效率）会下降，降低加速效果，就得不偿失了，所以我们提出伴生训练的方式，目的是为了更新mhead；
在精度无损的要求下，backbone部分需要实时更新，而mhead的更新可以延后于主干（只是接收率下降），且medusa的训练在RL的上游sft阶段已经完成训练，进入到RL之后再次进行微调，此时的训练数据已经变为rollout-batch-size，所以只需要按照rollout-bs的大小进行少量微调，所以只需要单台机器即可训练medusa，并且伴生训练可以和训练阶段overlap，参数更新的耗时也是无损的；
性能收益来看，medusa会有一定的计算开销，batchsize增大会降低性能收益，但即便是running-batch打满的情况下，仍然有50%性能收益。
4.1 通用能力提升
我们用开源模型RL后可以获得通用能力提升，这里列了两个case：
第一个是指令跟随，比如prompt为午餐吃什么，原始sft也是正常语义，但回答的是早餐，RL之后可以正确回答午餐内容；
第二个case为丰富性，比如prompt提问是否能够解决内部冲突问题，原始sft也能正确回答，但是回答太过简短，只有一个是的，相比之下RL的回答还有额外的解释，当然丰富性通过response_len也可以看出RL之后的长度更长；
整体上，基于开源模型做RL后，相比基线sft GSB提升5-20%，不同模型有一定差异，模型越强增益越小；右图所示为某开源模型在RL之后的各项任务打分，推理、数学提升较多，综合自动化评分有6%的提升，表明RL 能够提升通用能力。
4.2 PRM效果提升
我们也开发了PRM也就是基于过程的训练模式，基于PRM的过程奖励，可以增强可解释性，比如这个case是推理小赵换到小吴位置需要几个步骤，标准答案是需要换3次，PRM不仅推理结果正确，而且中间推理过程也能对不合理的部分进行惩罚，比如中间错误的说辞以及冗余的推断逻辑。
通过PRM能够让激励更准确，比如在相同的输出情况下，ORM和PRM的分差reward-gap会加大，更有区分度，更细力度的打分也能指导推理的过程，最终提升模型的上限5%。
4.3 调参经验
•可视化逐样本/token进行细粒度分析
•Advantage-batch对于小DP场景较为有用，防止走的偏
•Critic从RM进行参数加载，复用训练集群进行reward的batch-forward提高推理效率
•LR：基于sft进行设置，actor<critic，否则容易发散< span="">
•Critic先学习，actor先frozen，更准确的学习
•reward hacking奖励攻击：奖励后期的收敛性vs评测结果的差异。
PPO调参需要依赖一个鲁棒的系统，一些可视化的分析是有必要的，除了常用的tensorboard看loss/reward训练指标之外，还可以通过上图所示的具体rollout信息，进行逐样本更细粒度的分析；
调参方面有一些细节，比如advantage平滑的时候在整个rollout-batch进行，尽可能增大激励的稳定性，防止走偏；critic从RM进行参数加载，走本地offload可以保证精度的同时提高Reward的forward效率；lr一般基于sft进行设置，actor要小于critic，否则容易发散；尽可能让critic先学习，frozen住actor，让critic具备一定的评判能力，保证走的更准确；还有reward-hacking问题，训练指标可能持续增加，但评测集已经开始下降一般就是reward-hacking了，需要进行rm的迭代训练。
4.4 未来规划
未来计划主要围绕两个主题，一个是进一步的性能优化，一个是算法探索：
性能优化方面，MATRIX（MultimodalAI Training and Inference eXecutor）是AGI团队进行多模态大语言模型大规模分布式预训练、可监督微调、人类偏好对齐训练和在线推理的统一框架；
比如结合内部自研的推理引擎进一步压缩generate的耗时，不同seq的负载均衡，目前Pipeline优化主要针对rollout阶段，后续也可以结合训练做深度流水调度。
AGI的算法上，其实预训练的scaling已经吃的差不多了，通过post-train还有一定空间，单一的sft无法满足，通过RL可以进一步提升，甚至是infer阶段的scaling，比如现在OpenAI的O1实现，提升推理性能，虽没有公开细节，但已经有一些揣测，比如按照RL的迭代流程，sft是无法推断正确，但通过PRM的RL训练后，模型有了更强的推理能力，也是对RL的scaling-law的探索。
唯世
小红书大模型分布式训练工程师，负责多模态分布式训练框架的研发和优化。
云开
小红书大模型分布式训练工程师，负责多模态分布式训练框架的研发和优化。
霍金
小红书大模型分布式训练工程师，负责多模态分布式训练框架的研发和优化。
杜弼
小红书大模型分布式训练工程师，负责多模态分布式训练框架的研发和优化。
