在人工智能技术的快速发展中，多模态大语言模型（MLLM）以其强大的图文理解、创作、知识推理及指令遵循能力，成为了推动数字化转型的重要力量。然而，如何使这些模型的输出更加贴近人类的风格、符合人类的偏好，甚至与人类价值观保持一致，成为了一个亟待解决的问题。为了应对这一挑战，基于人类反馈信号的强化学习方法（RLHF）应运而生，其中，PPO（Proximal Policy Optimization）算法作为 OpenAI 的核心技术，在 RLHF 阶段扮演着关键角色。

小红书大模型团队，在这个技术日新月异的时代，开始了他们自研 MLLM RLHF 训练框架的征程。他们深知，要构建一个高效、准确的 RLHF 训练系统，需要综合考虑算法优化、系统架构、训练调度以及推理引擎等多个方面。

该图为我们的整体架构，上面是make-experience的第一步骤，dataloader迭代器返回prompt之后，通过vllm的generate输出response，prompt和response组合之后成为新的forward数据集；左下方是第二步骤，也就是进行4个模型的forward，生成最终的ppo-elements或者简称rollout，构造最终的train数据集；右方为PPO训练，actor训练之后会进行参数同步，这里会涉及到medusa-sft的协同训练流程，最终同步给vLLM进行新一轮的generate等。

首先是一些常规的训练性能优化方法，数据加载部分采用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并发数，把整体的吞吐进一步提升。

进一步分析发现，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等上下文开销，在开源模型上训练时性能最优。

针对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。

在大规模的场景下，图文混合场景增大的复杂度，也带来了新的瓶颈，因此需要进行特定的性能优化：

我们观察到视觉部分计算量大，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，能进一步降低预处理的耗时。

训推一致性是对工程架构的一个考验，比如训练之后的模型走推理框架评估可能会存在打分差异，特别的，算法上严格要求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的调优选项对齐等细节处理。

Rollout阶段的generate耗时始终是一个大头，所以我们结合推理部分进一步优化，但传统量化等方式有一定的精度损失，为了稳定PPO训练，也要要求无损优化。

投机采样是一种无损优化，通过多个头的预测来降低decode次数，因为自回归部分主要为 Memory-bound，进行计算并行优化，也是一种计算换空间的方式。这里我们采用的是medusa算法，由主干网络加上多个medusa-head组成，每个MedusaHead是和LM-Head一样结构，分别对应到预测每个头的topk的token，并且RL场景下需要组batch关注的是吞吐，不需要考虑延迟，所以running-batch尽可能大才能提高整体吞吐量；

但RL场景下存在一个问题，actor模型在不断更新，medusa的接收率（猜测准确越高，降低decode次数越多，越能提升推理效率）会下降，降低加速效果，就得不偿失了，所以我们提出伴生训练的方式，目的是为了更新mhead；

在精度无损的要求下，backbone部分需要实时更新，而mhead的更新可以延后于主干（只是接收率下降），且medusa的训练在RL的上游sft阶段已经完成训练，进入到RL之后再次进行微调，此时的训练数据已经变为rollout-batch-size，所以只需要按照rollout-bs的大小进行少量微调，所以只需要单台机器即可训练medusa，并且伴生训练可以和训练阶段overlap，参数更新的耗时也是无损的；

性能收益来看，medusa会有一定的计算开销，batchsize增大会降低性能收益，但即便是running-batch打满的情况下，仍然有50%性能收益。

我们用开源模型RL后可以获得通用能力提升，这里列了两个case：

第一个是指令跟随，比如prompt为午餐吃什么，原始sft也是正常语义，但回答的是早餐，RL之后可以正确回答午餐内容；

第二个case为丰富性，比如prompt提问是否能够解决内部冲突问题，原始sft也能正确回答，但是回答太过简短，只有一个是的，相比之下RL的回答还有额外的解释，当然丰富性通过response_len也可以看出RL之后的长度更长；

整体上，基于开源模型做RL后，相比基线sft GSB提升5-20%，不同模型有一定差异，模型越强增益越小；右图所示为某开源模型在RL之后的各项任务打分，推理、数学提升较多，综合自动化评分有6%的提升，表明RL 能够提升通用能力。

我们也开发了PRM也就是基于过程的训练模式，基于PRM的过程奖励，可以增强可解释性，比如这个case是推理小赵换到小吴位置需要几个步骤，标准答案是需要换3次，PRM不仅推理结果正确，而且中间推理过程也能对不合理的部分进行惩罚，比如中间错误的说辞以及冗余的推断逻辑。

通过PRM能够让激励更准确，比如在相同的输出情况下，ORM和PRM的分差reward-gap会加大，更有区分度，更细力度的打分也能指导推理的过程，最终提升模型的上限5%。

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的迭代训练。

未来计划主要围绕两个主题，一个是进一步的性能优化，一个是算法探索：

性能优化方面，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的探索。

作者介绍