浅析主流 Alignment 算法与 NeMo-Aligner 框架

今年 10 月在费城开 COLM 的时候，我有幸见到了 NVDA 两篇工作的 post，一篇是 RULER，现在已经是 long context 几乎必测的 benchmark；另一篇便是今天要讨论的工作——NeMo-Aligner。

https://github.com/NVIDIA/RULER
https://github.com/NVIDIA/NeMo-Aligner

Aligner 这个名字自然是非常恰当，毕竟“Aligner 并不试图贡献新的 Alignment 算法，而是专注于如何集成更多的 Alignment 算法”。所以，我姑且就用 Aligner 这一名字称呼这一系列的工作：

1. 
   

1. DeepSpeed-Chat——微软；

2. NeMo-Aligner——NVDA；

3. OpenRLHF——开源社区；

4. veRL——字节；

https://github.com/microsoft/DeepSpeedExamples/blob/master/applications/DeepSpeed-Chat/README.md
https://github.com/NVIDIA/NeMo-Aligner
https://github.com/OpenRLHF/OpenRLHF
https://github.com/volcengine/veRL

虽然 Aligner 是为了各类基于 RL 的 Alignment training 而生的，然而这些算法会有相当一部分时间用于推理（rollout），所以推理引擎社区也乐于关注 Aligner 工作的进展和需求。在下文中，我也会尽力从我的认知水平出发，讨论 Aligner 对于 Inference Engine 的需求。

Introduction

1. Nemo-Aligner 是一套集成了主流对齐算法的 toolkit，涵盖 RLHF / DPO / SteerLM 和 SPIN。能够有效调度千卡规模的计算资源，完成 Llama 4.1 405B 这种规模模型的训练。此外，toolkit 也支持 PEFT。作者希望这一框架足够 extenable，不过考虑到 RLHF / DPO 的差距已经非常大了，能复用的部分有限。个人感觉想要 support 一个新的算法还是比较困难。

2. 当前的 Alignment 算法基本还是地主的赛场，需要显著的计算资源。譬如 PPO 算法的计算流中，需要同时有四个模型进行复杂的交互。如果用 405B 的模型同时充当这四个模型的 base，在不进行优化的情况下，动用的资源可以简单算一算。

70B 模型用 FP16：140GB
4 个彼此交互的模型：140GB * 4 = 560GB

Adam：每个参数需要 8 个字节（两个动量）：70B * 8bytes = 560GB
Critic 和 Policy 同时都需要训练，所以需要两个优化器 = 1120GB

激活值和梯度：保守估计是模型大小的 1.5 倍，大概 210GB

KL 散度和其他各类开销：保守估计 200GB

大概 2100GB，也即 3 台 80G A100 ？？？

3. 基于我不严谨的计算，可以发现运行这样的 aligner 框架的开销非常恐怖。据此，Nemo-Aligner 需要在系统上下大功夫。首先，在 Megatron-LM 上 continue，支持 distributed 3D (data, tensor, pipeline) parallelism training。再者，在 rollout（可以理解为 sampling 阶段）使用 TensorRT-LLM 来做 inference，毕竟目前 trt-LLM 的推理效率确实是 SOTA。（而且也是英伟达自家的产品 ）

Model Alignment Background

SFT

pre-trained / base model 通过 Supervised Fine-Tuning 的方式来迫使模型基于 cross-entropy mimic 人类期望的回答。SFT 对于 DPO / RLHF 是必须的，因为不做 SFT，模型几乎无法 follow 人类指令。这一过程也被称为 behavior cloning。

RLHF

1. 这是今天的主角。RLHF 避免了显式为 RL 定义 reward function，转而通过 reward model 给出的 judgement 作为 reward。reward 模型从一组 pairwise dataset 中训练得到，一个 instruct 给出两个 response，其中一个标为 chosen，另一个标为 rejected。reward model 的 loss 基于 Bradley-Terry 模型得到，尝试让 Reward(chosen) > Reward(reject) 的概率尽可能大。RLHF 有两大主流方法，一种是 REINFOCE（1992），相对古老；而 NeMo-Aligner 更倾向于 PPO（2022）。

2. 下图一定程度表示了 PPO 的过程，这里尤其需要注意到模型参数是否会被训练。事实上，如前文所述，SFT Model、Policy / Actor、Reward Model、Value / Critic 这四个模型都会在计算流中被使用。其中 SFT 和 Reward 是 inference only 的，而 Policy 和 Value 会被更新。

3. SFT 已经论述过了，简单提一下 Reward Model Training。一般而言，reward model 会在 SFT model 基础上加上一层 linear layer，然后连同这个 linear layer 训练，将 linear layer 的 final project value 作为 reward。（在 SGLang 中体现为 classify 接口）

https://sgl-project.github.io/backend/native_api.html#Classify-(reward-model)

4. 计算流可以如下概括：

1. 初始化：SFT model 初始化 Actor，Reward Model 初始化 Critic；

2. Rollout / 生成输出：Actor 对输入的 Prompts 生成相应 / responses。

3. Reward Compute / 计算奖励：Reward Model 对生成的 responses 进行评分，提供 reward 得分。

4. 价值评估与更新：Critic 计算当前 responses 的价值（注意是 value 而非 reward），然后基于 value
   和 reward 计算优势函数（advantage function），最终确定 Actor 的更新方向和幅度。此外，Critic
   也会在这一步进行更新，减少其预测的 value 和实际奖励之间的误差，从而提高下个状态 value 的精度。

5. Advantage funciton 的值最红会用于优化 Actor，同时这步更新需要基于 SFT Model 的 logits 计算
   KL 散度并且加以惩戒，避免 Actor 相较于 SFT model 偏离太远，走到了 Reward Model 的盲点 / blind
   spots 上。

DPO

1. DPO 是和 PPO 同样声名大噪的对齐方法，这是一种 offline, off-policy 的对齐方法。试图直接将模型对齐成为符合偏好的最佳 policy 而避开了显式需要一个 reward model。这也是其名字的由来：Your Language Model is Secretly a Reward Model。（PS：最近发现有好几个给论文起名的新 pattern，譬如 XXX is All You Need，XXX is Secretly a YYY，XXX Empowers XXX，最后还有 Make XXX Great Again ）

2. 相较于使用隐式的 reward model，DPO 更为激进地选择使用 reference policy 通过 Bradley-Terry Mdel 隐式推导出 chosen / reject pair 之间的 reward。具体来说，通过计算 chosen/ reject pair 之间的对数概率差值来得到 reward，而该差值又由 optimal and reference policies 共同计算。这样的差值又通过 scaling 并且经过 sigmoid 得到最终的 loss。训练过程中，reference policy 固定不动，仅用于构造 chosen / rejected responses。

SteerLM

1. SteerLM 和 DPO 类似，避免复杂的 RL 方式，仅仅是 supervised finetuning。

2. 计算流如下：

1. 训练属性属性预测模型 / Attribute Prediction Model（APM），用于给定一个输入 prompt，给出多个语义评分。
   譬如对 correctness、toxicity 分别打 0 ~ 4 分。

2. 利用 APM 标标注 prompt-response pair，得到属性分。

3. 做 SFT：输入是 prompt + 属性分字符串，模型学习对给定 prompt 和属性分目标生成相应的 response。

4. 推理的时候，在 prompt 结尾加上需要的属性分即可。

3. 听上去真的是比其他方法简单多了，避免了 Reward Model 和 PPO 训练的复杂性，将属性对齐问题转为了显式的条件对齐。不过，猜测效果不是那么好

Self-Play Fine-Tuning

这是我们组的工作 ，不由得说影响力真不错。SPIN 和 DPO 类似，绕开了显式使用 reward model（注意到 reward model 的意义在于绕开了显式使用 reward function ）。SPIN 方法中，strong model 会从一个 weak model 的自我博弈中迭代而来。具体而言，给定一个 prompt / response pairs，让 weak model 对 prompt 进行一次新的标注，得到 prompt / (response, generated response) 这样的 preference data。然后，policy 在这样自我合成了一般数据的 preference pair 上训练，使其对于给定的 prompt，其偏向 response 的概率高于偏向先前自我构造的 generated response。这里需要用到 preference loss function，而这样的 preference loss function 和 DPO 中使用的 loss function 一模一样。

Online / Offline vs On-policy / Off-policy

作为强化学习白纸，这里记录下自己理解的概念。

1. On-Policy vs Off-Policy：用于收集经验的策略（behavior policy）和被优化的策略（target policy）是否是同一个。 On-Policy 使用当前正在学习的 policy 收集数据，数据只能使用一次，用后即弃，更加稳定容易收敛；但是数据利用率低，构造数据成本高昂，典型算法有 PPO 和 SARSA。而 Off-Policy 可以使用任意的 policy 来构造数据，可以复用历史数据，可以使用外部策略（如人类）的数据，数据利用率高，但是训练可能不稳定，需要复杂的重要性采样，典型算法如 Q-Learning，DQN，SAC。

2 . Online vs Offline：算法是否需要在训练过程中与环境交互。 Online 算法在训练过程中不断与环境交互，实时收集数据，动态调整策略，可以更适应环境变化，然而实时访问的风险和成本较高。典型场景有实时控制、在线决策。Offline 算法又称为 Batch RL， 只使用预先收集的数据集 ，训练过程不与环境交互，类似于监督学习。安全，不需要实时环境，但是受限于数据质量，难以处理分布外的情况。

3. PPO 是 online on-policy 的；但是对于 Online 而言，我和组里的同学讨论，这个定义可能没那么严格。Online 的特点可能是小批量大轮次，比如 SPIN 中一轮 1w 个 preference data 然后 train 3 个 epoch，而典型的 PPO 可能是 64 个 preference data train 成百上千轮，所以 SPIN 可能不是典型的 Online。

4. 【from 我们组的学长】 从 RL 的角度来看，是否从当前 policy 采样决定是否是 on policy；而是否和环境交互决定是否是 on line。这种界限并不清晰，没必要纠结，知道在描述什么即可。

NeMo-Aigner For RLHF / PPO

1. 讨论了非常久的 alignment background，现在我们回到 NeMo-Aligner 上。之后的解析（对 OpenRLHF 和 veRL）应该不会涉及 RL 背景了。RLHF 的流程如本文第一张图所示，而主要的开销还是在 PPO 流程中，具体而言是 rollout。

2. 如前文所述， PPO 过程需要 4 个模型组件参与整个流程。

a. PPO Actor：训练且推理，由 SFT Model 初始化而来，是 PPO 最终希望得到的微调结果。

b. Reference Policy：仅推理，一般设置为 SFT Model 本身，不进行任何修改。在 PPO 过程中，Actor 会

Reference Policy 同时计算 logits【这里会在后文论述】，然后计算 KL 散度，防止模型过度偏移。

c. PPO Critic：训练且推理，从 reward model 初始化而来，在 PPO 中计算 value【并非 reward】。

d. Reward Model：仅推理，对 Actor 产生的 rollout data 提供 rewards。

3. 上述四种 model 都可以任意大小，因此 NeMo-Aligner 实现了分布式训练，通过 PyTriton 启动 server 和 client，从而不再要求 critic 和 actor 在同一个节点上。

4. 直观上，可以启动四个模型 server 来同时 host 四个组件。然而，注意到 PPO Actor 和 Reference Policy 实际上是同一架构同一参数量的模型，不过参数会有所更新。因此，NeMo-Aligner 将二者组合在同一个 job 上。二者不同时使用， 所以可以将其中不用的一个 offload 到 DRAM 上，推理时再异步 swap 回来。 【这里也有很深的文章可以做，后文分析】这种 offload and swap 的策略在 PPO Critic 和 Reward Model 上也同样适用。

5. Response generation 占据了 PPO 全流程的主要时间。显然，我们可以采用 inference engine 来加速这样的 generation。为什么不可以用 training engine 的 forward 来进行 generation 呢？这个问题问的看似愚蠢，仔细想想，为什么我们没有用 training engine 来做 inference，比如直接拿着 deepspeed 的 forward 来做推理，反而要单独设计推理引擎呢？其实还是推理的计算特性所致。

https://zhuanlan.zhihu.com/p/4148050391

6. 在之前的文章（上方链接到的文章）中有提到，inference（或者说占据主要时间的 decode）是 autoregressvie 的，每个 token 依赖于前序所有 token，所以无法单个 sequence 的 decode 是无法并行的（能组起的 batch 都是多个 sequence 的 continuous batching），从而计算密度并不高。另一方面，decode 阶段需要进行大量的内存读取和数据传输，对通讯要求较高。 总归，inference engine 是主要被 memory bounded 的。 而 training engine 则不然， training engine 的 batch size 可以开的非常大，直到打满 engine 的计算能力，从而使 compute bounded 的。 因此，用 training engine 来做 inference 并不科学。（尽管我的本科写过的作业从不考虑这个问题 ）

7. 如上所述，我们已经发现，对于 PPO 中的四个 components，我们分别需要 inference engine 和 trainning engine，彼此优化各自的计算目标。现在考虑这个问题：为了计算 KL 散度需要得到 reference policy 和 actor 分别的 logits，这个 logits 该由什么 engine 得到？显然，inference engine 得到 logits 快很多，而且现代引擎都支持这个请求。 然而，目前 inference engine 得到的 logits 精度是更低的，不应该来计算 KL 散度。更科学的方法是用 trainning engine 得到二者的 logits 来计算 KL 散度。

8. 至于为什么 inference engine 的 logits 精度更低，这也是为了 inference 速度做出的牺牲。直观上来说，continuous batching 组的 batch 越大，精度”飘“的更厉害。即便 batch size 写死为 1，也不够准确。

9. 考虑到 engine 启动的开销很大，PPO 过程中任意 engine 都不应该被关掉。这无疑带来了更大的显存压力，为此 NeMo-Aligner 在反向传播时重新计算 training 阶段的激活值，减少了峰值显存压力。此外，inference 需要的显存小于 training，所以在单个节点的显存允许时，NeMo-Aligner 在 inference 时只采用 tensor parallelism，避免了 pipeline parallelism 带来的跨节点通讯开销。

10. 在随后的迭代训练过程中，inference 的 engine 需要与 training engine 更新过的参数进行同步。【这里有很大文章】 为了实现这样的同步，inference engine 使用了 Tensor RT Refitter 来进行 in-place update，而不是关掉 engine 重启新的