从0到1构建RLHF系统——小红书大模型团队的探索与实践

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

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

1.1 RL强化学习原理

在介绍本文工作之前，我们先了解一下RL强化学习的原理。

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进行设计和优化的；

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 整体架构

下面介绍我们的整体系统架构，框架选型上我们吸取了开源比如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 异构组网架构

下面介绍异构组网架构，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 同构组网架构

接下来介绍同构组网架构，当数据量、模型规模等都增大的情况下，集群规模压力也大幅增大，我们以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，比如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进行采样阶段，训练无需进行，存在资源浪费，即便是通过训推资源的混合切换，也会存在较大的空窗；因此我们设计了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阶段的整体耗时。

我们整体对比一下不同的实现方式：

1. 常规版本：每个节点独立运行各自的任务，推理在推理集群上部署，ac的训练和推理也在各自的集群上部署，但是推理服务和训练性能差异较大，存在一定优化空间；

2. 全量offload：相当于vllm/actor/critic分别进行全集群的资源利用，这种方法实现简单，可以充分提高generate的并发能力，roi较高；但是也引入了明显的offload/restart等空隙；

3. 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优化

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

我们观察到视觉部分计算量大，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 训推一致性

训推一致性是对工程架构的一个考验，比如训练之后的模型走推理框架评估可能会存在打分差异，特别的，算法上严格要求RM（Reward-Model）的精度一致，最早通过RM-serving进行accuracy统计发现掉点，经过排查发现是因为框架的diff，逐位误差在千分位，但是RM的vhead是一个linear，本质上是一个矩阵乘，相当于扩大了hidden平方倍，精度无法容忍；

我们的解决方法如下：

1. RM+RL的训练均使用相同的训练框架；

2. RL场景下要把RM-serving走api的方式转为本地offload去做forward，这样能够保证训推完全一致，且没有额外的serving占用和通信开销；

3. 训推一致性，训练和推理采样mcore的计算workload，保证评测任务上训推tokenwise一致，如下图所示，多个评测case上完全相同；这里也会涉及到一些技术实现，比如通过allgather+local-reduce实现替换allreduce，cublas的调优选项对齐等细节处理。

3.6 Medusa提升采样效率

投机采样是一种无损优化，通过多个头的预测来降低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%性能收益。

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效果提升

4.3 调参经验

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的探索。

唯世

小红书大模型分布式训练工程师，负责多模态分布式训练框架的研发和优化。

云开

小红书大模型分布式训练工程师，负责多模态分布式训练框架的研发和优化。

霍金

小红书大模型分布式训练工程师，负责多模态分布式训练框架的研发和优化。

杜弼

小红书大模型分布式训练工程师，负责多模态分布式训练框架的研发和优化。