小红书大模型团队RLHF系统从零构建实践

在人工智能技术快速迭代的当下，多模态大语言模型（MLLM）展现出的图文理解、创作、知识推理和指令遵循能力，已经成为推动数字化转型的关键力量。但问题也随之而来——如何让模型的输出更贴近人类风格？如何让它们遵循人类偏好，甚至与人类价值观保持一致？

基于人类反馈的强化学习方法（RLHF）正是为此而生。其中，OpenAI 的核心技术 PPO（Proximal Policy Optimization）算法在 RLHF 阶段扮演着至关重要的角色。

小红书大模型团队在这个技术日新月异的时代，启动了自研 MLLM RLHF 训练框架的征程。构建一个高效、准确的 RLHF 训练系统，需要综合考虑算法优化、系统架构、训练调度以及推理引擎等多个维度——这绝不是一项简单的任务。

1.1 RL强化学习原理

在进入正题之前，先快速回顾一下RL强化学习的基本原理。

状态S： 输入的 prompt
动作A： 输出的 response
奖励R： 根据 prompt+response 由奖励模型打分

给定prompt，调整策略模型（policy），目标是生成符合人类喜好的response——这个喜好由RM（Reward-Model）提供的偏序信号来定义。

RL强化学习在游戏、自动驾驶等场景中应用较多。简单来说，就是智能体在特定环境中尝试不同策略的动作，并根据反馈不断更新状态——这是一个自闭环的过程。对应到NLP领域，大语言模型就是智能体，输入的prompt对应状态S，输出的下一个token对应动作A，对prompt+response的打分对应奖励R。整体目标就是：给定 prompt，调整 policy 策略模型，使其生成符合人类偏好的输出。

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四个模型的协同训练和推理，设计和实现一套高效准确的RLHF训练系统，确实是一项关键挑战。

PPO训练流程

PPO是一个训推混合的训练流程，可以分为两个步骤：make-experience（经验采样）和training（训练）。

第一阶段是经验采样，有时也叫rollout。在RL中，这就像一次实验——actor模型在当前环境下执行策略动作。具体实现是：actor根据prompt数据集生成response，然后基于prompt+response进行前向计算，得到logp、values、reward等元素。这个阶段涉及actor、ref、critic、reward四个模型的推理过程。

第二阶段是训练流程，涉及actor和critic两个模型。从损失函数计算来看，两个模型算法上是相互独立的，本质上是各自独立训练。

工程实现的痛点是显而易见的：涉及多阶段的dataloader（prompt、rollout、train），actor模型的generate自回归过程，还有四个不同的模型推理和两个不同的模型训练。

不过，make-experience阶段给优化留下了空间——接下来看看我们是怎么做的。

2.1 整体架构

架构选型

训练方面，我们从apex、nemo、megatron-lm等多个候选方案中比较，最终选择了Megatron-core做训练、ray做调度的组合；推理方面则采用vLLM推理引擎。

为什么这么选？之前的开源方案如trlx、openrlhf，往往采用apex、nemo、megatron等混合框架的组合拼凑，代码维护成本和更新负担太重。降维到megatron-core+ray的组合，最大的好处就是降低维护成本。推理服务使用vLLM，不过需要注意的是，vLLM虽然也能输出训练需要的logp数据，但为了与训练保持一致，我们坚持走megatron框架——任何框架间的差异都可能成为训练的bias，引入不必要的风险。

整体架构可以这样理解：上方是make-experience的第一步，dataloader迭代器返回prompt后，通过vLLM的generate输出response，prompt和response组合后成为新的forward数据集；左下方是第二步，对四个模型进行前向计算，最终生成PPO的rollout，构造train数据集；右方是PPO训练，actor训练后会进行参数同步，这里涉及medusa-sft的协同训练流程，最终同步给vLLM进行新一轮的generate。

2.2 异构组网架构

从0到1优化设计

1. Forward-Offload： 纯串行实现，分时复用模型，从4个模型降为2个。
2. Training-Seperated： actor和critic在独立集群上异步训练（并行化）。
启动流程： 在actor集群上用torchrun启动主进程训练，并自动以mock方式拉起critic所在集群——通过remote-func实现。

RL的workload涉及四个模型，如果全部串行实例化每个模型，成本相当高。但根据前面的分析，这些模型结构是相同的——于是我们设计了异构组网架构来大幅提升性能。

在Forward阶段，四个模型通过offload方式串行运行——actor完成后卸载，加载ref；critic完成后卸载，加载RM。这样模型需要的显存规模从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 同构组网架构

新的瓶颈出现，见招拆招

异构组网的痛点逐渐显现：

1. 数据量、序列长度、模型参数量增加，集群规模压力持续加大。以llama3 70B为例，SFT模型本身需要4机32卡，而在32k长文下开启CP2就需要8机64卡。
2. rollout同步问题：序列长度和多模数据量增加，同步的通信耗时也随之增加。

AC同构组网架构： 将actor和critic复用同一集群，main-model从2个降为1个——即tripe-off。
forward阶段：通过offload参数进行模型切换；
training阶段：offload中间状态（包括param、grad、ddp-bucket、opt-state等）。

当数据量和模型规模持续增大，集群规模压力也随之增大。还是以llama3 70B为例，原始异构组网需要16机128卡——资源压力不小。此外异构组网需要master-worker架构拉远实现，需要同步rollout的buffer数据，当多模数据增加图片等数据后，会大幅增加rollout同步的开销。

于是我们进一步实现了同构组网。前面提到forward阶段两个模型已经同构，训练阶段虽然是异步的，但actor和critic部分也可以进一步通过offload完成scale。最终集群规模可以降低到单个SFT训练所需的规模。具体实现上，forward阶段只需offload参数就能完成模型切换；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引擎
• vLLM优化： 调整切分方案，虽然降低running-batch，但能提高并发

先看一些常规的训练性能优化方法。数据加载环节采用prefetch，在ptx-loss等场景下使用双dataloader时，需要降低num-workers数来防止CPU OOM。并行策略上，使用megatron的TP、PP、CP、SP并行方法。在长文本和大模型场景下，采用recompute方式进行显存优化。dynamic-batch方面，训练和推理使用不同的batchsize——因为推理不需要梯度优化器等显存占用，所以batchsize可以更大。vLLM部分采用round-robin方式将prompt请求分发到多个vLLM引擎，实现负载均衡。推理服务的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资源合并到训练中
• 流水线优化：最大化make-experience阶段的效率

进一步分析发现，actor进入采样阶段时，训练流程其实处于闲置状态，集群资源明显有浪费。那么能不能把这些空窗期利用起来？我们设计了一套pipeline方案，核心思路是训推集群混布，增大serving并发，压缩generate耗时。具体实现上，vLLM和megatron都支持全量offload——vLLM部分支持kvcache销毁和param offload，forward部分涉及actor/critic的param和其他buffer的offload切换。最终可以把vLLM资源合并到训练内部，实现资源最大化。再通过类比PP并行，把generate和forward阶段做成流水线并行，对mini-batch进行细粒度切分，大幅降低make-experience阶段的整体耗时。

三种实现方式的对比：

1. 常规版本： 每个节点独立运行各自任务，推理在推理集群部署，actor和critic的训练推理各自在不同集群。推理服务和训练性能差异较大，存在明显优化空间。
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优化

MLLM（Multimodal Large Language Model）

痛点： VLM部分需要图像处理，图文混合导致负载不均衡。

特点：

1. 视觉部分计算量大，但参数量较小
2. PPO场景下图像处理存在冗余

优化方案：

1. 组网优化：LLM的TP/PP/CP用作视觉部分的DP
2. 多路复用：vLLM/actor/critic复用相同的img_feature
3. Prefetch：冻结visual-model，通过prefetch机制，融入make-experience做Pipeline的overlap

在大规模场景下，图文混合带来了额外的复杂度，也出现了新的瓶颈。我们观察到，视觉部分计算量大，但LLM部分模型权重更大。针对这个现象，可以对视觉部分和LLM使用不同的并行策略：视觉部分只用DP，LLM部分用MP、PP、CP的组合。将LLM部分的PP、MP、CP用作视觉部分的DP，充分利用资源。

MLLM针对特定训练任务进行优化，目前SFT和RL阶段都采用冻结visual-model的方式，只微调adapter部分。因此可以进行多路复用——vLLM、actor、critic三者共享相同的img_feature。同时通过prefetch机制，把vit视觉部分的整个实现放到dataloader内部，与make-experience阶段结合做Pipeline overlap，进一步降低预处理耗时。

3.5 训推一致性

确定性训练和推理

问题：

1. RM的精度要求——RM-serving时出现accuracy掉点
2. 推理评估的精度要求

原因： RM的vhead是一个linear结构，本质上是矩阵乘——[1, hidden] * [hidden, 1]。如果logits存在框架层面的差异，误差就会被放大：初始diff约0.001，矩阵乘后可能达到0.3。

训推一致性是对工程架构的严峻考验。训练之后的模型走推理框架评估，可能会有打分差异。特别地，算法上对RM（Reward-Model）的精度一致性有严格要求。最早通过RM-serving进行accuracy统计时发现了掉点，排查后发现是框架差异导致的——逐位误差在千分位级别，但RM的vhead是linear结构，矩阵乘相当于放大了hidden倍的平方，精度完全无法接受。

解决方法如下：

1. RM和RL的训练使用同一套训练框架
2. RL场景下，将RM-serving通过API的方式转为本地offload做forward——这样能保证训推完全一致，且没有额外的serving占用和通信开销
3. 训推一致性方面，训练和推理采用mcore的计算workload，确保评测任务上训推tokenwise一致。这里涉及一些技术实现细节：比如用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一样，分别对应预测每个头的top-k token。在RL场景下，由于是组batch运行，重点是吞吐而不是延迟，所以running-batch尽可能大才能提高整体吞吐。

伴生训练设计

问题： Backbone和mhead的匹配问题——接收率下降。

伴生训练：

1. 精度无损：backbone实时更新，mhead也同步更新
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训练后，通用能力有明显提升。这里列举两个案例：

第一个是指令跟随能力。比如输入"午餐吃什么"，原始SFT模型也能理解语义，但回答内容是关于早餐的。经过RL之后，模型能正确输出午餐相关的回答。

第二个是输出丰富性。比如询问"能否解决内部冲突问题"，原始SFT模型虽然回答正确，但回答过于简短——只有一个"是"。相比之下，RL后的回答增加了额外的解释和扩展。通过response length也能看出，RL后的输出明显更长。

整体来看，基于开源模型做RL后，相比基线SFT，GSB指标提升在5-20%之间。不同模型有差异——模型越强，增益越小。右图展示了某个开源模型在RL后的各项任务打分：推理和数学能力提升较多，综合自动化评分提高约6%。这说明RL确实能提升通用能力。

4.2 PRM效果提升

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

通过PRM，奖励信号更准确。同样输出结果下，ORM和PRM的分差（reward-gap）会更大，区分度更好。更细粒度的打分能指导推理过程，最终模型上限提升约5%。

4.3 调参经验

• 可视化：逐样本和逐token进行细粒度分析
• Advantage-batch：在小型DP场景下特别有用，防止训练走偏
• Critic从RM加载参数，复用训练集群进行reward的batch-forward，提高推理效率
• LR：基于SFT设置，actor的学习率要小于critic，否则容易发散
• 让critic先学习、actor先冻结，确保critic具备评判能力，走得更稳
• reward hacking（奖励攻击）：注意奖励后期的收敛性与评测结果之间的差异

PPO调参需要一个鲁棒的支撑系统。除了常用的tensorboard观察loss和reward训练指标外，还可以通过上图中的具体rollout信息进行更细粒度的逐样本分析。

调参的细节比较多：比如advantage平滑时在整个rollout-batch范围内进行，尽可能增大激励的稳定性，防止走偏；critic从RM加载参数，走本地offload既能保证精度，又能提高reward的forward效率；学习率一般基于SFT设置，但actor要小于critic，否则容易发散；尽可能让critic先学习，冻结actor，让critic先具备评判能力，保证后续训练更准确。还要注意reward-hacking问题：训练指标可能持续上升，但评测集已经开始下降——这通常就是reward-hacking的标志，需要重新迭代训练RM。

4.4 未来规划

未来计划围绕两个主题：性能优化和算法探索。

性能优化方面，MATRIX（Multimodal AI Training and Inference eXecutor）是AGI团队进行多模态大语言模型大规模分布式预训练、可监督微调、人类偏好对齐训练和在线推理的统一框架。后续计划包括：结合内部自研的推理引擎进一步压缩generate耗时；处理不同序列长度下的负载均衡；目前Pipeline优化主要针对rollout阶段，未来可以结合训练做深度流水调度。

在AGI算法层面，预训练的scaling空间已经比较有限，post-train还有提升空间——单一的SFT无法满足需求，RL可以进一步释放能力，甚至延伸到推理阶段的scaling。比如OpenAI的O1实现，通过RL方式提升推理性能——虽然细节未公开，但已有不少推测。按照RL的迭代流程，SFT无法正确推断的问题，通过PRM的RL训练后，模型获得了更强的推理能力。这本身就是对RL scaling-law的有效探索。

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

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

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

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

你是一名 AI 行业编辑，请围绕下面这条热点输出一份资讯解读：
热点：小红书大模型团队RLHF系统从零构建实践要求：
1. 先用一句话解释这条热点在讲什么
2. 再总结它为什么重要
3. 说明会影响哪些 AI 产品或内容方向
4. 最后给出 3 个适合资讯站使用的标题