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）需要在训练期间运行一个复杂的流程，涉及四个大型语言模型以复杂的方式交互。这种对齐算法引入了新的系统挑战，需要重新思考软件堆栈的各个方面，包括模型可扩展性、模型之间的协调以及训练循环中的文本生成。
现有的开源模型对齐工具，最著名的是HuggingFace TRL（von Werra et al., 2020）、CarperAI triX（Havrilla et al., 2023）和Microsoft DeepSpeedChat（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），试图最大化 $r_{chosen}>r_{rejected}$ 的可能性（即预测的奖励与人类偏好一致）。一旦Reward Model训练完成，它就可以用于计算RL算法的奖励。RLHF中最常用的两种方法是REINFORCE（Williams, 1992）和近端策略优化（PPO）（Schulman et al., 2017）。在NeMo-Aligner中，我们专注于PPO，特别是Ouyang et al. (2022) 描述的PPO。
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训练了Llama3（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：Llama3 模型的并行设置
| 模型 | 计算节点数量 | 张量并行 | 流水线并行 | 数据并行 | 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，得益于框架设计。
8 结论
现代模型对齐技术，尤其是基于强化学习的技术，在系统实现方面提出了复杂的优化挑战。我们创建并开源了NeMo-Aligner，以允许AI研究人员和从业者通过以可扩展的方式利用所有可用计算来高效地进行LLM对齐实验。我们的框架在训练大型模型时始终表现出良好的扩展性。由于这是我们的初始版本，我们预计随着未来版本的发布，这种扩展性只会进一步提高。此外，我们支持在计算受限的设置中使用LoRA进行SFT、PPO、DPO、SteerLM的参数高效微调。作为Apache 2.0许可的开源代码库，NeMo-Aligner可以使对齐研究更加高效和易于访问。
致谢
我们要感谢NVIDIA的许多团队为NeMo-Aligner的启用做出的贡献，特别是NeMo、TRT-LLM和TensorRT团队。