RLHF 训练框架 NeMo-Aligner
摘要
将大型语言模型（LLMs）与人类价值观和偏好对齐对于使其变得有用且安全至关重要。然而，构建高效的对齐工具具有挑战性，尤其是对于包含数百亿参数的最大、最强大的 LLMs。我们创建了 NeMo-Aligner，这是一个模型对齐工具包，可以高效扩展到数千个 GPU，用于训练最大的开源 LLMs，如 Nemotron 4 340B 和 Llama 3.1 405B。NeMo-Aligner 提供了高度优化和可扩展的实现，支持主要的模型对齐范式，如：基于人类反馈的强化学习（RLHF）、直接偏好优化（DPO）、SteerLM 和 Self-Play Fine-Tuning （SPIN）。此外，我们的工具包支持在参数高效微调（PEFT）设置下运行大多数对齐技术。NeMo-Aligner 设计为可扩展，允许以最小的工作量支持其他对齐技术。它以 Apache 2.0 许可证开源，我们邀请社区贡献，地址为 https://github.com/NVIDIA/NeMo-Aligner。
1. 引言
在大量未标记文本上预训练大型语言模型展示了其强大的能力（Brown et al., 2020; Zhang et al., 2022）。虽然这种无监督预训练模型取得了令人印象深刻的成果，但随后将模型对齐以遵循用户指令是挖掘 LLMs 在实际应用中的潜力的关键步骤（Sanh et al., 2022; Wei et al., 2022）。基于监督微调（SFT）的尝试（Conover et al., 2023; Kopf et al., 2023; Taori et al., 2023）被证明不如那些使用反馈来调整模型以生成更有帮助的响应并避免生成不太有用的响应的技术有效（Bai et al., ; Ouyang et al., 2022; Touvron et al., 2023; Dong et al., 2023）。
尽管使用反馈训练模型有诸多好处，但这些流程非常难以正确实施（Lambert & Calandra, 2023; Zheng et al., ），阻碍了其在少数资源充足的组织之外的广泛、高效采用。例如，基于人类反馈的强化学习（RLHF）方法中的近端策略优化（PPO）变体（Ouyang et al., 2022）需要在训练期间运行一个复杂的流程，涉及四个大型语言模型以复杂的方式交互。这种对齐算法引入了新的系统挑战，需要重新思考软件堆栈的各个方面，包括模型可扩展性、模型之间的协调以及训练循环中的文本生成。
现有的开源模型对齐工具，最著名的是 Hugging Face TRL （von Werra et al., 2020）、CarperAI triX （Havrilla et al., 2023）和 Microsoft DeepSpeed-Chat （Yao et al., 2023）。这些工具在可用性和功能集方面提供了极好的起点。然而，通过 NeMo-Aligner，我们的目标是大幅提高性能和可扩展性，支持超过一千个 GPU，特别适用于对齐最大、最强大的模型，如 Nemotron 4 340B （Nvidia et al., 2024）、Llama 3.1 405B （Dubey et al., 2024）及更高版本。
NeMo-Aligner 通过以下方式解决可扩展性挑战：
（I）基于 Megatron-LM （Shoeybi et al., 2020）的 3D （数据、张量和流水线）并行训练
（II）采用分布式方法进行 RLHF 中的近端策略优化（PPO）训练
（III）在 rollout 阶段集成基于 TensorRT-LLM （NVIDIA, ）的 PPO 推理优化
这些优化相结合，使用户能够高效地在数千个 GPU 上训练最大模型，减少研究迭代时间。
NeMo-Aligner 优化了流行的对齐技术，包括监督微调（SFT）、基于 PPO 的 RLHF （Ouyang et al., 2022）、直接偏好优化（Rafailov et al., 2023）、SteerLM （Dong et al., 2023）和 Self-Play Fine-Tuning （Chen et al., 2024）。我们在第 2 节简要概述了这些技术的背景，随后在第 3、4、5 和 6 节中深入探讨了每种技术的训练。最后，我们在第 7 节展示了 NeMo-Aligner 的可扩展设计。
2. 模型对齐背景
图 1：基于 Ouyang et al. (2022) 的 RLHF 训练流程。 步骤 1：使用带注释的提示-响应数据对预训练（基础）模型进行监督微调（SFT）。 步骤 2：使用偏好数据训练 SFT 模型以生成 Reward Model。 步骤 3： SFT 模型用于初始化策略网络， Reward Model 用于初始化价值网络——结合输入提示，所有四个模型用于训练策略模型。SFT 模型还用于计算步骤 3 中的 KL 散度惩罚（未图示）。
2.1 监督微调
给定一个预训练（也称为“基础”）模型，监督微调（SFT）在带有预期响应的提示上更新基础模型的参数，其中预期响应可能来自专家人工注释（Kopf et al., 2023）或其他语言模型（Ding et al., 2023）。模型被训练为在给定提示的情况下模仿预期响应，使用基于标记的交叉熵损失。SFT 是基于人类反馈的强化学习（Ouyang et al., 2022）和直接偏好优化（Rafailov et al., 2023）的重要前提步骤，因为没有它，基础模型不太可能生成遵循用户指令的响应。此步骤有时也称为行为克隆，因为模型被期望模仿人类或其他模型的响应。
2.2 基于人类反馈的强化学习
基于人类反馈的强化学习（RLHF）由 Christiano et al. (2017) 提出，旨在避免在强化学习中使用手动定义的奖励函数。相反， Reward Model 是从由“选择”和“拒绝”轨迹对组成的人类偏好数据集中训练的。Reward Model 的损失源自 Bradley-Terry 模型（Bradley & Terry, 1952），试图最大化
RLHF 已被证明在模型对齐方面带来了显著的好处（Ouyang et al., 2022; Bai et al., ; Touvron et al., 2023），其典型的训练流程如下，如图 1 所示：
- 从预训练的基础模型开始，训练一个初始的 SFT 模型，如第 2.1 节所述。
- 从 SFT 模型开始，使用由“选择”和“拒绝”响应组成的人类偏好数据集训练 Reward Model，遵循 Christiano et al. (2017) 的方法。通常，我们在 SFT 模型之上初始化一个线性 Reward Model 头，然后进行训练。
- 从 SFT 模型开始，使用在线近端策略优化算法（PPO, Schulman et al., 2017）训练策略，奖励由训练好的 Reward Model 提供。输入提示不一定与用于 Reward Model 训练的提示相同。基于 SFT 模型的 KL 散度的正则化项有助于防止策略偏离其起点太远，并避免利用 Reward Model 的“盲点”（Stiennon et al., 2020; Ouyang et al., 2022）。PPOCritic 通常从 Reward Model 初始化。
2.3 直接偏好优化
直接偏好优化（DPO）（Rafailov et al., 2023）是一种离线、离策略算法，它利用偏好数据直接训练最优策略，而无需显式的 Reward Model。与使用 Reward Model 不同， DPO 使用 Reference Policy 通过 Bradley-Terry 模型隐式推导出“选择”和“拒绝”对之间的奖励。这是通过计算最优策略和 Reference Policy 之间的对数概率差异来实现的，该差异被缩放并通过 sigmoid 函数转换为损失。Reference Policy 在训练期间被冻结，并代表用于生成给定提示的“选择”/“拒绝”响应的策略。如果用于生成偏好数据的 Reference Policy 不可用，可以通过在偏好数据的提示和首选响应上进行监督微调来近似。
2.4 SteerLM
SteerLM （Dong et al., 2023）是一种基于监督微调的模型对齐算法，类似于 DPO，避免了复杂的 RL 方法。SteerLM 涉及三个步骤。
第一步是训练一个属性预测模型，该模型学习预测响应的各种语义方面的值（在 0 到 4 之间，值越高越好），这些方面使响应更有帮助且更安全，例如其正确性和毒性（Kopf et al., 2023; Wang et al., 2023）。
接下来，属性预测模型可用于注释多样化的提示-响应数据集中有助于帮助性和安全性的各种属性。
最后，这些注释的数据集可用于执行属性条件监督微调，其中模型学习生成响应，条件是提示以及格式化为字符串的注释属性，例如 helpfulness:4, correctness:4, toxicity:0。此步骤教会模型以细粒度的方式区分更有帮助/更安全的响应和不太有帮助/不太安全的响应。在推理时，提示可以附加最佳属性值，如上所述，以生成最有帮助的响应。
2.5 Self-Play Fine-Tuning
Self-Play Fine-Tuning （SPIN）（Chen et al., 2024）是一种基于自对弈的算法，其中通过与前几代的自身对弈，从较弱的模型发展出更强的模型。从提示/响应对的 SFT 数据集开始，从前几代模型生成新的响应。然后，通过区分这些自生成的响应和地面真实的人类生成的 SFT 响应来改进其策略。这是通过一个与 DPO （第 2.3 节）相同的偏好损失函数来实现的。当 SPIN 训练开始时，我们使用初始策略的副本作为 DPO 损失中的 Reference Policy。然后，自对弈“游戏”进行多次迭代，在此期间我们像 DPO 一样训练策略，同时保持 Reference Policy 冻结，并在每次迭代结束时用训练策略的权重更新 Reference Policy 的权重。在每次迭代中，我们遍历 SFT 训练数据集，并使用 Reference Policy 为每个提示生成响应，构建地面真实 SFT 人类“选择”响应和生成的“拒绝”响应之间的偏好元组。一旦我们为整个 epoch 构建了这些偏好元组，我们就通过这些“（选择，拒绝）”偏好对的 DPO 损失函数更新模型权重。因此，模型隐式地学习偏好地面真实 SFT 响应，而不是前几代自身生成的响应，这构成了自对弈机制。
3. RLHF （PPO）训练
NeMo-Aligner 旨在高效地支持多种对齐技术，并在极大规模上运行。它通过基于 Megatron-LM （Shoeybi et al., 2020）和 NeMo （Kuchaiev et al., 2019）构建，包含诸如 Transformer Engine （NVIDIA, 2022）的优化内核、分布式融合 Adam 优化器和 3D 并行支持等功能。NeMo-Aligner 支持 Ouyang et al. (2022) 引入的整个 RLHF 流程，如第 2.2 节所述。训练流程分为三个不同的阶段，如图 1 所示：监督微调、Reward Model 训练和近端策略优化。该流程的效率挑战主要来自近端策略优化阶段，本节描述了我们在应对这些挑战时的方法，如图 2 所示。
图 2： RLHF 训练的优化。PPO 训练和推理的优化分别在第 3.1 节和第 3.2 节中详细说明。
3.1 PPO 训练的分布式方法
PPO 阶段需要在四个不同的模型上运行训练和/或推理，如图 3 所示：
- PPO Actor（训练和推理，从 SFT 模型初始化）：我们希望通过 PPO 微调的模型。
- Reference Policy（仅推理，设置为 SFT 模型）：用于计算 KL 惩罚的模型。
- PPO Critic（训练和推理，从 Reward Model 初始化）：用于在 PPO 中计算价值估计的模型。
- Reward Model（仅推理）：为生成的 rollout 数据提供 RL 奖励。
所有这些模型都可能非常大（例如 Llama 3.1 405B），因此 NeMo-Aligner 采用分布式方法进行 PPO 训练。我们允许用户设置 PyTriton （NVIDIA, 2022）服务器和客户端，以便在 PPO 期间跨不同模型进行通信。这些 PyTriton 服务器使得可以在不同的计算集群上运行模型，消除了将 Critic 和 Actor 放在同一计算分配上的要求。最初，会启动四个不同的服务器（即每个模型一个）。然而，我们注意到 Reference Policy 和 PPO Actor 是相同的模型，但具有不同的权重。因此，我们将它们组合到一个作业中，并将 Reference Policy 的权重卸载到 CPU，在 Reference Policy 推理步骤中与 Actor 的权重交换。我们对 Reward Model 和 Critic 采用相同的策略。所有通信都是异步进行的，允许 Critic 推理/训练与策略推理/训练的流水线操作。
我们扩展计算分配大小，使得[Reward Model 推理 + Critic 推理] ≈ [Actor 采样 + Reference Policy 推理]，并且[Critic 训练] ≤ [Actor 训练 + Actor 推理初始化]。这确保了流水线可以最有效地利用可用计算容量。
图 3： NeMo-Aligner PPO 系统架构。PPO Actor 是一个 PyTriton （NVIDIA, 2022）客户端，它向服务器（PPO critic 和奖励模型）发送异步请求，以获取生成 rollout 的奖励和价值，并将训练数据发送给 critic。
3.2 PPO rollout 的优化
在 rollout 步骤中的响应生成主导了 PPO 训练的端到端时间。Actor 的生成阶段由多个前向传递组成，每次前向传递生成一个 token。因此，生成阶段的内核通常受启动延迟和内存带宽的限制，这意味着直接重用训练阶段的计算优化前向传递实现会导致非常差的性能。
为了解决这些瓶颈，我们使用 TensorRT-LLM （NVIDIA, ）实现生成阶段，这是一个高性能的 LLM 部署框架。TensorRT-LLM 将推理优化内核和自动内核融合集成到基于 TensorRT 的运行时中，以实现更好的性能。在 RLHF 开始时，模型被传递给 TensorRT-LLM，后者将模型编译为 TensorRT 引擎； TensorRT-LLM 将引擎加载到其运行时中并执行生成。引擎保存了模型权重的副本以及运行时 KV 缓存和激活。由于序列化引擎的成本，我们在训练期间将引擎保留在内存中。因此，我们通过在反向传递中重新计算训练阶段的激活来减少峰值内存压力。此外，由于生成的内存需求低于训练，如果内存允许，我们在推理期间仅使用张量并行对模型进行重新分片，从而消除在流水线并行运行时节点间通信的开销。
在后续的训练步骤中，必须将引擎与训练阶段的更新参数权重同步。引擎使用 TensorRT Refitter （NVIDIA, ）就地更新。我们避免重新编译引擎，因为这会带来很大的开销，因为生成在权重更新之前无法开始。
对于较大的问题规模，由于响应长度的差异，生成时间在最快和最慢的数据并行工作器之间可能存在差异。为了缓解这种情况，我们允许用户设置工作器池，以动态平衡数据并行工作器之间的负载，为生成较短响应的工作器分配相应更多的工作。
3.3 模型训练细节和质量
作为使用 NeMo-Aligner 进行大规模 RLHF 训练的实践演示，我们使用 PPO 训练了 Llama 3 （Meta AI, 2024） 70B 模型，如 Wang et al. (2024) 所述。PPO 模型使用 128 的 rollout 全局批量大小、128 的训练全局批量大小、1e-7 的恒定学习率、0.003 的 KL 惩罚进行训练，并使用 HelpSteer2 作为提示源。
根据 Wang et al. (2024); Meng et al. (2024)，我们使用 MT-Bench （Zheng et al., ）和 GPT-4-Turbo 评判器来评估训练后的 RLHF 模型的性能。最终模型在 MT-Bench 上取得了 8.13 的成绩，相较于初始的 SFT 检查点的 7.96 有所提升。Reward Model 和 RLHF 训练后的模型已公开发布，地址分别为 https://huggingface.co/nvidia/llama3-70B-SteerLM-RM 和 https://huggingface.co/nvidia/llama3-70B-PPO-Chat。
3.4 可扩展性
表 1：在不同计算节点数量下扩展训练的效果（Llama 3 70B Actor 和 Llama 3 70B critic， rollout 全局批量大小为 128， BF16 精度）
| 计算节点数量（Actor + Critic） | 8 + 4 | 16 + 8 |
|---|---|---|
| 每步时间（秒，标准差）↓ | ||
| 总体 | 53.7 (2.223) | 29.8 (1.459) |
| 训练 | 8.5 (0.356) | 4.8 (0.118) |
| Rollout | 45.2 (1.957) | 25.0 (1.41) |
| - 响应生成 | 35.8 (1.658) | 17.5 (1.224) |
| - Log-probs 计算 | 4.0 (0.573) | 2.8 (0.275) |
| - TensorRT Refit | 3.1 (0.182) | 2.3 (0.334) |
| - Critic 等待 | 0.01 (0.001) | 0.01 (0.001) |
| 相对加速（vs. 8 + 4 节点设置）↑ | ||
| 总体 | 1x | 1.80x |
| 训练 | 1x | 1.77x |
| Rollout | 1x | 1.81x |
| - 响应生成 | 1x | 2.05x |
| - Log-probs 计算 | 1x | 1.43x |
为了展示 NeMo-Aligner 的扩展效率，我们重复了第 3.3 节中的相同训练设置，分别使用 8 个 Actor 节点 + 4 个 critic 节点和 16 个 Actor 节点 + 8 个 critic 节点。如表 1 所示，每步的总时间相应减少，在 8+4 节点和 16+8 节点之间实现了 1.80 倍的加速。每步总时间的加速得益于训练和 rollout 阶段的加速，展示了 NeMo-Aligner 在这两个阶段的有效优化。
训练阶段的扩展是次线性的，因为随着节点数量的增加，每个数据并行 rank 的微批次数量减少。由于流水线并行模型中的所有流水线阶段必须在调用优化器之前完成，因此我们产生了填充和排空流水线的开销，这与微批次的数量无关（Shoeybi et al., 2020）。因此，减少每个数据并行 rank 的微批次数量会增加训练步骤中填充和排空流水线的时间比例，此时 GPU 利用率较低。在 log prob 计算阶段，扩展比例为 1.43x，也存在类似的问题。
表 2：扩展训练的效果（Llama 3 70B Actor 和 Llama 3 70B critic， rollout 全局批量大小为 1024， BF16 精度）
| 每步时间（秒，标准差）↓ | 16 + 8 | 32 + 16 | 64 + 32 |
|---|---|---|---|
| 总体 | 190.4 (15.392) | 106.8 (6.842) | 56.9 (7.596) |
| 训练 | 38.8 (4.674) | 22.2 (2.408) | 14.1 (1.439) |
| Rollout | 151.5 (13.172) | 84.6 (5.596) | 42.7 (6.36) |
| - 响应生成 | 131.3 (12.551) | 71.7 (5.093) | 29.9 (6.456) |
| - Log-probs 计算 | 15.5 (3.055) | 8.5 (1.753) | 6.2 (1.177) |
| - TensorRT Refit | 2.2 (0.408) | 2.2 (0.014) | 3.5 (0.063) |
| - Critic 等待 | 0.02 (0) | 0.03 (0.001) | 1.7 (2.045) |
| 相对加速（vs. 16 + 8 节点设置）↑ | |||
| 总体 | 1x | 1.78x | 3.35x |
| 训练 | 1x | 1.75x | 2.75x |
| Rollout | 1x | 1.79x | 3.55x |
| - 响应生成 | 1x | 1.83x | 4.39x |
| - Log-probs 计算 | 1x | 1.82x | 2.50x |
生成时间随着节点数量的增加而很好地扩展，实现了接近线性的 2.05 倍加速。这是因为增加 Actor 节点的数量会成比例地增加每个步骤的数据并行工作器数量，这些工作器可以均匀分担生成的高强度工作。然而，将权重拟合到 TensorRT 引擎中的时间相对恒定，因此无法很好地随节点数量扩展。最后， Actor 和 critic 模型之间的异步通信使得等待 critic 模型的额外时间可以忽略不计（0.01 秒），这表明在 PPO 流水线中使用 Actor 和 critic 模型之间的异步非阻塞调用的有效性。
表 3： Llama 3 模型的并行设置
| 模型 | 计算节点数量 | 张量并行 | 流水线并行 | 数据并行 | Rollout 批量大小 |
|---|---|---|---|---|---|
| 70B + 70B | 8 + 4 | 8 + 8 | 8 + 4 | 1 + 1 | 8 |
| 70B + 70B | 16 + 8 | 8 + 8 | 8 + 4 | 2 + 2 | 8 |
| 70B + 70B | 32 + 16 | 8 + 8 | 8 + 4 | 4 + 4 | 8 |
| 70B + 70B | 64 + 32 | 8 + 8 | 8 + 4 | 8 + 8 | 16 |
| 405B + 405B | 84 + 42 | 8 + 8 | 21 + 21 | 4 + 2 | 16 |
系统可扩展性还需要在问题需求的背景下考虑。第 3.3 节中的训练设置有一个 70B Llama 3 Actor、70B Llama 3 Critic 以及 128 的 rollout 全局批量大小。这样的设置限制了我们在超过 16+8 节点的情况下有效展示系统扩展，因为没有足够的工作可以有意义地分配给更多的数据并行工作器。因此，我们在表 2 中稍微修改了设置，使用 1024 的 rollout 全局批量大小，以在需求更高时测量系统性能。表 2 显示，训练作业的增加需求使其能够有意义地扩展到 64+32 节点（总共 768 个 H100 GPU）以进行 PPO 中的各个阶段。
3.5 什么贡献了系统性能？
为了更好地理解 NeMo-Aligner 的 PPO 系统设计中每个方面的重要性，我们通过逐一移除一个方面并测量每步的总时间来进行消融研究，如表 4 所示。我们发现， TensorRT-LLM 集成是系统高性能的最关键组件，没有它， PPO 每步将花费近七倍的时间。其次是仅在推理期间使用张量并行对模型进行重新分片（3.87x），使用 TensorRT Refit 避免 TensorRT-LLM 引擎重新编译（3.15x），以及使用 Actor 和 critic 模型之间的异步请求（1.54x）。我们没有观察到使用工作器池来平衡数据并行 rank 之间的工作有显著加速，因为问题规模较小（rollout 全局批量大小为 128），因此工作器不平衡小于我们进行平衡本身的开销。然而，我们预计此功能在问题规模较大时会有所帮助。
表 4：消融研究（Llama 3 70B Actor 和 critic， rollout 全局批量大小为 128， 8 个 Actor 节点和 4 个 critic 节点）
| 每步时间（秒，标准差）↓ | 相对于最佳 RLHF 设置的时间↓ | |
|---|---|---|
| 最佳 RLHF 设置 | 53.7 (2.223) | 1x |
| - TensorRT-LLM 集成 | 372.2 (37.433) | 6.93x |
| （即使用 NeMo Generate） | ||
| - Reshard | 207.6 (8.940) | 3.87x |
| - TensorRT Refit | 169.0 (5.206) | 3.15x |
| - 异步请求 | 82.8 (7.826) | 1.54x |
3.6 训练最大的开源 LLM
表 5： NeMo Aligner 支持最大的开源 LLM （Llama 3.1 405B）
| 计算节点数量（Actor + Critic） | 84 + 42 |
|---|---|
| 每步时间（秒，标准差）↓ | |
| 总体 | 164.60 (3.434) |
| 训练 | 5.60 (0.532) |
| Rollout | 158.90 (2.969) |
| - 响应生成 | 140.10 (1.067) |
| - Log-probs 计算 | 17.20 (2.352) |
| - TensorRT Refit | 0.70 (0.108) |
| - Critic 等待 | 0.30 (0.454) |
NeMo-Aligner 支持截至 2024 年 7 月最大的开源 LLM，如 Nemotron 4 340B （Nvidia et al., 2024）和 Llama 3.1 405B （Dubey et al., 2024）。在表 5 中，我们使用 1008 个 H100 GPU 对 Llama 3.1 405B 进行 PPO，使用基于第 3.3 节的配置。我们使用 Llama 3.1 405B Instruct 作为 Actor，并在 Llama 3.1 405B Instruct 上训练 Reward Model，使用与 Nemotron 4 340B Reward （Wang et al., 2024）相同的数据和超参数。与 70B 模型相比， 405B 模型在 PPO 训练期间明显更慢，主要瓶颈在于缓慢的响应生成阶段。
这是因为 405B 模型的生成无法适应单个节点，需要流水线并行生成，我们计划在未来的工作中进一步优化。此外，在 PPO 的前几步中， 405B 模型生成的响应比 70B 模型更长（平均 916 个 token，而 70B 模型平均为 351 个 token）。
4. DPO 训练
我们遵循 Zephyr-7B-Beta （Tunstall et al., 2023）的训练配方，该模型使用 SFT 和 DPO 进行训练。简而言之，首先在 Mistral-7B （Jiang et al., 2023）上使用 Ultrachat 数据集（Ding et al., 2023）进行 SFT。然后使用 Ultrafeedback 数据集（Cui et al., 2023）进一步进行 DPO 训练。对于 SFT，我们使用 2e-5 的恒定学习率、512 的全局批量大小，并训练模型 3 个 epoch。对于 DPO 训练，我们使用 3e-4 的 KL 正则化系数、512 的全局批量大小和余弦学习率计划，峰值 LR 为 1e-7，最小 LR 为 1e-8， 50 个预热步骤，最多 300 步。我们获得了比 Tunstall et al. (2023) 报告的略好的 MT-Bench 分数，无论是最终模型（7.60 vs 7.34）还是仅 SFT 的初始模型（6.77 vs 6.64）。
5. 使用 LoRA 进行 SteerLM 训练
低秩适应（LoRA）（Hu et al., 2021）使得以更高效和更具成本效益的方式微调大型语言模型成为可能。在 NeMo-Aligner 中支持各种对齐技术， LoRA 应用于 SteerLM 训练，遵循 Wang et al. (2023) 的训练配方，使用 Llama 2 70B 模型以及 HelpSteer （Wang et al., 2023）和 Open Assistant 数据集（Kopf et al., 2023）。具体来说，我们将 LoRA 应用于所有注意力层，秩为 32。我们使用 128 的全局批量大小、1e-5 的恒定学习率，在 10 个预热步骤后使用 AdamW 优化器，并训练 3 个 epoch。如表 6 所示，将 LoRA 应用于 SteerLM 训练，使用 BF16 可以将所需的 80GB GPU 的最小数量从 32 减少到 8。在相同数量的 GPU 下， LoRA 实现了 5 倍的加速，同时保持了可比的模型性能： MT-Bench 7.43 vs. 7.54，这在该基准测试的噪声水平内（Jiang et al., 2023）。
表 6：全参数和 LoRA SteerLM 的比较
| 全参数 | LoRA | |
|---|---|---|
| 可训练参数数量 | 70B | 89M |
| 所需 80GB GPU 最小数量 | 32 | 8 |
| 相对速度（样本/GPU/秒） | 1 × | 5 × |
| MT-Bench | 7.54 | 7.43 |
随着我们增加用于 LoRA 训练的 GPU 数量，相对吞吐量（以每秒样本数衡量）几乎成比例提高，如图 4 所示。这表明 NeMo-Aligner 可以有效地在大量 GPU 上分配和并行化工作负载，且开销和收益递减最小。
6. SPIN 训练
我们通过 SPIN 重新创建了 Zephyr-7B-Beta （Tunstall et al., 2023） SFT 模型，而不是 SFT，如 Chen et al. (2024) 所提出的。我们从 Mistral-7B 基础模型（Jiang et al., 2023）开始，并按照 Chen et al. (2024) 进行 SPIN 训练。然而，我们在他们的方法上做了一些调整，即我们没有将前几代的生成注入当前代（这会使每个 epoch 的数据集大小翻倍），并且我们只进行单次迭代，每次迭代 1 个 epoch。此外，我们仅使用 Ultrachat200k （Ding et al., 2023）中的 50k 样本的随机子集，而不是整个数据集，并使用 AdamW 而不是 RMSProp。我们的学习率为 5e-7，总共 400 步， 40 个预热步骤，然后使用余弦退火将 LR 衰减到 1e-7，持续最后 100 步。全局批量大小为 64，权重衰减为 0.0， KL 正则化系数为 0.1，如 Chen et al. (2024) 所述。使用这种方法，我们获得了 7.04 的 MT-Bench 分数，超过了使用 SFT 的 Zephyr-7B-Beta 的 6.64 （Tunstall et al., 2023），以及 3 次迭代 SPIN 模型的 6.78 （Chen et al., 2024）。
7. 框架可扩展性
我们设计 NeMo-Aligner 时考虑了可扩展性，允许用户轻松修改算法，尽管分布式训练复杂。我们通过使用训练器抽象来实现这一点，鼓励跨各种步骤和方法重用现有的训练器方法。NeMo-Aligner 的可扩展性允许以最少的代码更改集成 DPO 的变体，包括 Identity Preference Optimization （Azar et al., 2023）、Conservative DPO （Mitchell, 2023）和 Kahneman-Tversky Optimization （Ethayarajh et al., 2023）。此外，其他模型对齐技术，如 Constitutional AI （Bai et al., ）、Rejection Sampling （Touvron et al., 2023）和 Self-Rewarding Language Models （Yuan et al., 2024）也被纳入 NeMo-Aligner，得益于框架设计。
结论
现代模型对齐技术，尤其是基于强化学习的技术，在系统实现方面提出了复杂的优化挑战。我们创建并开源了 NeMo-Aligner，以允许 AI 研究人员和从业者通过以可扩展的方式利用所有可用计算来高效地进行 LLM 对齐实验。我们的框架在训练大型模型时始终表现出良好的扩展性。由于这是我们的初始版本，我们预计随着未来版本的发布，这种扩展性只会进一步提高。此外，我们支持在计算受限的设置中使用 LoRA 进行 SFT、PPO、DPO、SteerLM 的参数高效微调。作为 Apache 2.0 许可的开源代码库， NeMo-Aligner 可以使对齐研究更加高效和易于访问。
致谢
我们要感谢 NVIDIA 的许多团队为 NeMo-Aligner 的启用做出的贡献，特别是 NeMo、TRT-LLM 和 TensorRT 团队。