标题：NeMo-Aligner：用于高效模型对齐的可扩展工具包
作者：Gerald Shen, Zhilin Wang, Olivier Delalleau, Jiaqi Zeng, Yi Dong, Daniel Egert, Shengyang Sun, Jimmy Zhang, Sahil Jain, Ali Taghibakhshi, Markel Sanz Ausin, Ashwath Aithal, Oleksii Kuchaiev
机构：NVIDIA

A1 主要贡献

本文介绍了NeMo-Aligner，一个旨在高效地将大型语言模型（LLMs）与人类价值观和偏好对齐的工具包。对齐是使LLMs变得有用和安全的关键步骤，但为拥有数百亿甚至数千亿参数的大模型构建高效的对齐工具充满挑战。

核心问题：现有的模型对齐开源工具（如HuggingFace TRL、trlX、DeepSpeedChat）虽然为社区提供了良好的起点，但在性能和可扩展性方面，尤其是在对齐如Nemotron 4 340B和Llama 3.1 405B等顶级大模型时，面临瓶颈。复杂的对齐算法（如RLHF中的PPO）引入了新的系统挑战，包括模型扩展性、模型间协调以及在训练循环中进行文本生成等。

研究目标与创新点：NeMo-Aligner旨在大幅提升模型对齐的性能和可扩展性，使其能够高效扩展至上千块GPU，从而支持对最大规模的开源LLMs进行对齐。其主要贡献和创新点如下：
* 极致的可扩展性：通过以下三项关键技术应对扩展性挑战：
1. 基于Megatron-LM【33, Megatron-lm: Training multi-billion parameter language models using model parallelism, 2020, Arxiv】构建，支持3D（数据、张量、流水线）并行训练。
2. 为RLHF中的近端策略优化（PPO）训练实现了分布式方法。
3. 在PPO的rollout阶段集成了基于TensorRT-LLM【26, TensorRT-LLM, 2023, GitHub】的推理优化。
* 全面的对齐算法支持：提供了对主流模型对齐范式的高度优化和可扩展实现，包括：
* 基于人类反馈的强化学习（RLHF），特别是PPO变体。
* 直接偏好优化（DPO）。
* SteerLM。
* 自博弈微调（SPIN）。
* 高效的微调支持：支持在参数高效微调（PEFT）设置下运行大多数对齐技术，如LoRA。
* 框架的可扩展性：NeMo-Aligner被设计为可扩展的，用户能够以最小的努力支持其他对齐技术。
* 开源贡献：该工具包以Apache 2.0许可证开源，鼓励社区参与贡献。

A3 背景知识

本节概述了NeMo-Aligner支持的几种关键模型对齐技术。

图1：基于Ouyang et al. (2022) 的RLHF训练流程。步骤1：使用标注的Prompt-Response数据对预训练（基础）模型进行监督微调。步骤2：将得到的SFT模型用偏好数据进行训练，以产生一个奖励模型。步骤3：SFT模型用于初始化策略网络，奖励模型用于初始化价值网络——连同输入提示，所有四个模型都用于训练一个策略模型。SFT模型也用于在步骤3中计算KL散度惩罚（图中未示出）。

2.1 监督微调（Supervised Fine Tuning）

基本定义。给定一个预训练（或称“基础”）模型，监督微调（SFT）在包含提示和期望响应的数据集上更新模型参数。这些期望响应可能来自人类专家标注【18, Openassistant conversations – democratizing large language model alignment, 2023】或由其他语言模型生成【10, Enhancing chat language models by scaling high-quality instructional conversations, 2023】。模型通过词元级别的交叉熵损失函数进行训练，以模仿给定的期望响应。SFT是RLHF【30, Training language models to follow instructions with human feedback, 2022】和DPO【31, Direct preference optimization: Your language model is secretly a reward model, 2023】中一个重要的先决步骤，因为没有SFT，基础模型很难生成遵循用户指令的响应。这个步骤有时也被称为行为克隆，因为它期望模型模仿人类或其他模型的响应。

2.2 基于人类反馈的强化学习（Reinforcement Learning from Human Feedback）

RLHF概念与流程。RLHF由Christiano等人【7, Deep reinforcement learning from human preferences, 2017】提出，旨在避免在强化学习中手动定义奖励函数。它通过一个在人类偏好数据集（包含“选中”和“拒绝”轨迹对）上训练的奖励模型来实现。该奖励模型的损失函数源自Bradley-Terry模型【4, The rank analysis of incomplete block designs — I. The method of paired comparisons, 1952】，旨在最大化r_chosen > r_rejected的概率。奖励模型训练好后，可用于为RL算法计算奖励。RLHF中最常用的两种方法是REINFORCE【44, Simple statistical gradient-following algorithms for connectionist reinforcement learning, 1992】和近端策略优化（PPO）【32, Proximal policy optimization algorithms, 2017】。NeMo-Aligner专注于Ouyang等人【30, Training language models to follow instructions with human feedback, 2022】描述的PPO方法。RLHF已被证明对模型对齐有显著益处【30, Training language models to follow instructions with human feedback, 2022; 2, Training a helpful and harmless assistant with reinforcement learning from human feedback, 2022a; 37, Llama 2: Open foundation and fine-tuned chat models, 2023】，其典型训练流程如Figure 1所示，包含三个步骤：
1. SFT模型训练：从一个预训练的基础模型开始，如2.1节所述训练一个初始的SFT模型。
2. 奖励模型训练：基于SFT模型，使用一个包含“选中”和“拒绝”响应对的人类偏好数据集来训练一个奖励模型。通常，在训练前，我们在SFT模型之上初始化一个线性的奖励模型头。
3. PPO策略训练：从SFT模型出发，使用在线PPO算法训练一个策略模型，奖励由训练好的奖励模型提供。一个基于与SFT模型的KL散度的正则化项，用于防止策略偏离其起点太远并利用奖励模型的“盲点”【34, Learning to summarize from human feedback, 2020; 30, Training language models to follow instructions with human feedback, 2022】。PPO的critic（评论家）通常从奖励模型初始化。

2.3 直接偏好优化（Direct Preference Optimization）

DPO核心思想。DPO【31, Direct preference optimization: Your language model is secretly a reward model, 2023】是一种离线、离策略的算法，它利用偏好数据直接训练一个最优策略，而无需显式的奖励模型。它使用一个参考策略，通过Bradley-Terry模型来隐式地推导出“选中”和“拒绝”响应对之间的奖励。这是通过计算最优策略和参考策略下，“选中”和“拒绝”响应的对数概率之差来实现的。这个差值经过缩放并通过sigmoid函数转换，最终得到损失函数。在训练过程中，参考策略是冻结的，它代表了生成这些偏好数据（“选中”/“拒绝”响应）的策略。如果生成偏好数据的原始参考策略不可用，可以通过在偏好数据的提示和“选中”响应上进行监督微调来近似得到。

2.4 SteerLM

SteerLM方法。SteerLM【11, SteerLM: Attribute conditioned SFT as an (user-steerable) alternative to RLHF, 2023】是一种基于监督微调的模型对齐算法，与DPO类似，它避免了使用复杂的RL方法。SteerLM包含三个步骤：
1. 属性预测模型训练：首先，训练一个属性预测模型，该模型学习预测响应的各种语义属性的值（0到4之间，值越高越好），这些属性构成了有用和安全响应的要素，例如正确性和毒性【18, Openassistant conversations – democratizing large language model alignment, 2023; 41, Helpsteer: Multi-attribute helpfulness dataset for steerlm, 2023】。
2. 数据集标注：接下来，使用属性预测模型在多样化的提示-响应数据集中标注各种有助于提升“有用性”和“安全性”的属性。
3. 属性条件监督微调：最后，使用这些标注过的数据集进行属性条件监督微调。在这一步，模型学习在给定提示和格式化属性字符串（如helpfulness:4,correctness:4,toxicity:0）的条件下生成响应。这教会了模型以细粒度的方式区分哪些响应更有用/更安全。在推理时，可以将最优属性值附加到提示后，以生成最有帮助的响应。

2.5 自博弈微调（Self-Play Fine-Tuning）

SPIN机制。自博弈微调（SPIN）【6, Self-play finetuning converts weak language models to strong language models, 2024】是一种基于自博弈的算法，通过与自身先前版本的实例进行博弈，从一个较弱的模型发展出一个更强的模型。该过程从一个SFT数据集（提示/响应对）开始。首先，由模型的先前迭代版本生成新的响应。然后，通过区分这些自生成的响应和SFT数据中的真实人类响应来改进其策略。这是通过一个与DPO（2.3节）所用偏好损失函数相同的函数来实现的。当SPIN训练开始时，初始策略的一个副本被用作DPO损失中的参考策略。自博弈“游戏”会进行多轮迭代。在每轮迭代中，策略模型像DPO一样进行训练，而参考策略保持冻结；在每轮迭代结束时，用训练好的策略权重更新参考策略。在每次迭代中，算法遍历SFT训练数据集，并使用参考策略为每个提示生成响应，从而构建一个偏好元组，其中SFT中的真实人类响应为“选中”项，生成的响应为“拒绝”项。当整个epoch的偏好元组都构建好后，模型权重通过这些（“选中”，“拒绝”）对的DPO损失函数进行更新。通过这种自博弈机制，模型隐式地学会了偏好SFT的真实响应，而不是其先前迭代版本生成的响应。

A2 方法细节

NeMo-Aligner旨在高效地支持多种对齐技术，并在极大范围内实现扩展。它通过建立在Megatron-LM【33, Megatron-lm: Training multi-billion parameter language models using model parallelism, 2020, Arxiv】和NeMo【17, Nemo: a toolkit for building ai applications using neural modules, 2019】之上，集成了来自Transformer Engine的优化核【23, TransformerEngine, 2022, GitHub】、分布式融合Adam优化器以及3D并行支持等特性。NeMo-Aligner支持由Ouyang等人【30, Training language models to follow instructions with human feedback, 2022】引入并于2.2节描述的完整RLHF流程。该训练流程分为三个不同阶段（如图1所示）：监督微调、奖励模型训练和近端策略优化。流程效率的挑战主要来自PPO阶段，本节将描述我们应对这些挑战的方法，如图2所示。

图2：RLHF训练的优化。PPO训练和推理的优化分别在3.1节和3.2节中详述。

3.1 PPO训练的分布式方法

PPO阶段的模型构成与分布式架构。PPO阶段需要在四个不同的模型上运行训练和/或推理，如图3所示：
1. PPO Actor (训练和推理，从SFT模型初始化): 我们希望用PPO微调的模型。
2. Reference Policy (仅推理，设为SFT模型): 用于计算KL散度惩罚的模型。
3. PPO Critic (训练和推理，从奖励模型初始化): 在PPO中用于计算价值估计。
4. Reward Model (仅推理): 为生成的rollout数据提供RL奖励。

分布式实现与优化。所有这些模型都可能非常大（例如Llama 3.1 405B），因此NeMo-Aligner采用分布式方法进行PPO训练。我们允许用户设置PyTriton【22, PyTriton: Framework facilitating NVIDIA Triton inference server usage in Python environments, 2022】服务器和客户端，以便在PPO期间跨不同模型进行通信。这些PyTriton服务器使得模型可以运行在不同的计算集群上，消除了将critic和actor部署在同一计算资源分配上的要求。一种朴素的方法是为每个模型启动一个服务器（即四个服务器）。然而，我们注意到参考策略和PPO actor是相同的模型，只是权重不同。因此，我们将它们合并到一个作业中，并将参考策略的权重卸载到CPU，在进行参考策略推理步骤时与actor的权重进行交换。我们对奖励模型和critic也采用了相同的策略。所有通信都是异步完成的，允许critic的推理/训练与策略的推理/训练进行流水线操作。

计算资源平衡。我们调整计算资源的分配规模，使得 [奖励模型推理 + critic推理] ≈ [actor采样 + 参考策略推理] 并且 [critic训练] ≤ [actor训练 + actor推理初始化]。这确保了流水线能够最有效地利用可用的计算能力。

图3：NeMo-Aligner PPO系统架构。PPO Actor是一个PyTriton【22】客户端，它向服务器（PPO critic和奖励模型）发送异步请求，以获取生成rollouts的奖励和价值，并发送critic的训练数据。

3.2 PPO rollout的优化

Rollout阶段的性能瓶颈。在PPO训练的端到端时间中，rollout步骤中的响应生成占主导地位。actor的生成阶段由多个前向传播组成，每个前向传播生成一个token。因此，生成阶段的核通常受限于启动延迟和内存带宽，这意味着直接重用训练阶段计算优化的前向传播实现会导致性能非常差。

TensorRT-LLM集成优化。为了解决这些瓶颈，我们使用TensorRT-LLM【26, TensorRT-LLM, 2023, GitHub】（一个高性能的LLM部署框架）来实现生成阶段。TensorRT-LLM将推理优化的核和自动核融合集成到一个基于TensorRT的运行时中，以实现更好的性能。在RLHF开始时，模型被传递给TensorRT-LLM，后者将模型编译成一个TensorRT引擎；TensorRT-LLM将该引擎加载到其运行时并执行生成。该引擎持有一份模型权重以及运行时的KV缓存和激活值。由于序列化引擎的成本很高，我们在训练期间将引擎保留在内存中。因此，我们通过在反向传播中重新计算训练阶段的激活值来减少峰值内存压力。此外，由于生成比训练的内存需求低，如果内存允许，我们会重新分片模型，在推理期间仅使用张量并行，从而消除了在使用流水线并行时节点间通信的开销。

动态权重更新。在后续的训练步骤中，引擎必须与训练阶段更新的参数权重同步。我们使用TensorRT Refitter【27, TensorRT Refitter, 2023】对引擎进行原地更新。我们避免了重新编译引擎，因为这会产生巨大的开销，因为在权重更新完成之前无法开始生成。

负载均衡。对于大规模问题，由于响应长度的差异，不同数据并行工作节点在生成过程中的生成时间可能存在差异。为了缓解这个问题，我们允许用户设置一个工作池，在数据并行工作节点之间进行动态负载均衡，从而为生成时间较短的工作节点分配相应更多的工作。

3.3 模型训练细节与质量

Llama3 70B RLHF训练实践。作为使用NeMo-Aligner进行大规模RLHF训练的演示，我们按照Wang等人【42, Helpsteer2: Open-source dataset for training top-performing reward models, 2024】的指导，使用PPO训练了Llama3【21, Llama 3 model card, 2024】70B模型。该PPO模型的训练参数为：rollout全局批量大小128，训练全局批量大小128，恒定学习率1e-7，KL惩罚系数0.003，并使用HelpSteer2作为提示来源。

性能评估。我们遵循Wang等人【42】和Meng等人【20, Simpo: Simple preference optimization with a reference-free reward, 2024】的方法，使用MT-Bench【47, Judging llm-as-a-judge with mt-bench and chatbot arena, 2023a】和GPT-4-Turbo作为裁判来评估训练后的RLHF模型性能。最终模型在MT-Bench上取得了8.13分的成绩，相较于其强大的SFT起点（MT-Bench分数为7.96）有所提升。奖励模型和RLHF训练后的模型已公开发布。

3.4 可扩展性

70B模型扩展性测试。为了展示NeMo-Aligner的扩展效率，我们重复了3.3节中的相同训练设置，分别使用了8个actor节点+4个critic节点和16个actor节点+8个critic节点。如表1所示，每步的总体时间相应减少，在8+4节点和16+8节点之间实现了1.80倍的加速。这一加速得益于训练（Train）和Rollout两个阶段的速度提升，表明NeMo-Aligner对这两个阶段都进行了有效优化。

表1：在rollout全局批量大小为128和BF16精度下，将Llama 3 70B actor和Llama 3 70B critic的训练扩展到不同数量计算节点的效果，遵循3.3节配置。节点为8*H100-80GB-SXM，通过节点内NVLink【24】和节点间Infiniband【28】互连。每步时间是基于第一步之后5个步骤的平均值计算的，因为第一步为TRT-LLM引擎构建产生了大量时间。更多训练配置细节见表3。

扩展性分析。训练阶段的扩展是次线性的，因为随着节点数量增加，每个数据并行rank的微批量数量会减少。在流水线并行模型中，由于所有流水线阶段必须在优化器被调用前完成，我们会产生一个填充和排空流水线的开销，这个开销与微批量的数量无关【33, Megatron-lm: Training multi-billion parameter language models using model parallelism, 2020, Arxiv】。因此，减少每个数据并行rank的微批量数量会增加训练步骤中用于填充和排空流水线的时间比例，此时GPU利用率较低。类似的问题在log prob计算阶段也很明显，扩展比为1.43倍。生成时间的扩展性很好，达到了接近线性的2.05倍加速。这是因为增加actor节点数量会按比例增加每一步的数据并行工作节点数，这些节点可以均匀地分担密集的生成工作。然而，将权重载入TensorRT引擎所花费的时间相对恒定，因此随节点数量扩展不佳。最后，Actor和Critic模型之间的异步通信使得等待Critic模型的时间可以忽略不计（0.01秒），这表明在PPO流水线中，actor和critic模型之间使用异步非阻塞调用是有效的。

表2：在rollout全局批量大小为1024和BF16精度下，扩展Llama 3 70B actor和Llama 3 70B critic训练的效果。节点为8*H100-80GB-SXM，通过节点内NVLink【24】和节点间Infiniband【28】互连。每步时间是基于第一步之后5个步骤的平均值计算的，因为第一步为TRT-LLM引擎构建产生了大量额外时间。我们的分布式优化器特性使NeMo-Aligner在64+32节点配置中支持将rollout微批量大小从8翻倍到16。这显著加快了响应生成时间。更多配置细节见表3。

大规模问题下的扩展性。系统可扩展性还需要在问题需求的背景下考虑。3.3节的训练设置（70B Llama 3 Actor，70B Llama 3 Critic，rollout全局批量大小128）限制了系统扩展到16+8节点以上的有效演示，因为没有足够的工作量在更多数据并行工作节点间有意义地共享。因此，我们在表2中将rollout全局批量大小增加到1024，以衡量系统在更高需求下的性能。表2显示，增加的训练任务需求使其能够有意义地扩展到64+32个节点（共768个H100 GPU），并在PPO的各个阶段都表现出良好的扩展性。

表3：Llama3模型扩展性实验的并行设置。节点数和并行配置表示为 actor + critic。

3.5 是什么贡献了系统性能？

消融研究。为了更好地理解NeMo-Aligner的PPO系统设计中每个方面的重要性，我们进行了消融研究，即每次移除一个方面并测量每步的总体时间，如表4所示。我们发现，TensorRT-LLM集成是实现高系统性能最关键的组件，没有它，PPO每一步将花费近七倍的时间。其次是在推理期间将模型重新分片以仅使用张量并行（3.87倍），使用TensorRT Refit避免TensorRT-LLM引擎重新编译（3.15倍），以及在actor和critic模型之间使用异步请求（1.54倍）。我们没有观察到使用工作池来平衡数据并行rank之间的工作负载有明显的加速，因为问题规模较小（rollout全局批量大小为128），因此工作负载不平衡带来的影响小于我们进行平衡本身的开销。尽管如此，我们预计这个特性对于更大的问题规模会有帮助。

表4：在8个Actor节点和4个critic节点上，使用128的rollout全局批量大小训练Llama 3 70B actor和critic的消融研究。节点为8*H100-80GB-SXM，通过节点内NVLink【24】和节点间Infiniband【28】互连。每步时间是基于第一步之后5个步骤的平均值计算的，因为第一步为TRT-LLM引擎构建产生了大量额外时间。

3.6 训练最大的开源LLM

支持超大模型。NeMo-Aligner支持截至2024年7月最大的开源LLM的对齐，例如Nemotron 4 340B【29, Nemotron-4 340b technical report, 2024】和Llama 3.1 405B【12, The llama 3 herd of models, 2024】。在表5中，我们使用1008个H100 GPU对Llama 3.1 405B进行了PPO训练，配置基于3.3节。我们使用Llama 3.1 405B Instruct作为actor，并在其基础上训练了一个奖励模型，使用了与Nemotron 4 340B Reward【42, Helpsteer2: Open-source dataset for training top-performing reward models, 2024】相同的数据和超参数。

性能观察。与70B模型相比，405B模型在PPO期间的训练速度要慢得多，主要瓶颈在于缓慢的响应生成阶段。这是因为405B模型的生成无法装入单个节点，需要进行流水线并行的生成，我们计划在未来的工作中进一步优化。此外，在PPO的前几个步骤中，405B模型生成的响应比70B模型更长（平均916个token vs. 平均351个token）。

表5：NeMo Aligner支持截至2024年7月最大的开源LLM（Llama 3.1 405B）作为actor和critic。我们在rollout全局批量大小为128和BF16精度下运行模型。节点为8*H100-80GB-SXM，通过节点内NVLink【24】和节点间Infiniband【28】互连。每步时间是基于第一步之后5个步骤的平均值计算的，因为第一步为TRT-LLM引擎构建产生了大量额外时间。更多配置细节见表3。

4 DPO训练

复现Zephyr-7B-Beta。我们遵循了Zephyr-7B-Beta【39, Zephyr: Direct distillation of lm alignment, 2023】的训练流程，这是一个通过SFT和DPO训练的模型。简而言之，首先在Mistral-7B【16, Mistral 7b, 2023】上使用Ultrachat数据集【10, Enhancing chat language models by scaling high-quality instructional conversations, 2023】进行SFT。然后使用Ultrafeedback数据集【9, Ultrafeedback: Boosting language models with highquality feedback, 2023】通过DPO进一步训练模型。

训练配置与结果。对于SFT，我们使用恒定学习率2e-5，全局批量大小512，训练了3个epoch。对于DPO训练，我们使用KL正则化系数3e-4，全局批量大小512，以及一个余弦学习率调度，峰值学习率为1e-7，最低学习率为1e-8，50个预热步骤，最多300个步骤。我们获得的MT-Bench分数略高于Tunstall等人【39】报告的分数，最终模型为7.60（对比7.34），仅SFT的初始模型为6.77（对比6.64）。

5 使用LoRA进行SteerLM训练

LoRA在对齐中的应用。低秩自适应（LoRA）【15, Lora: Low-rank adaptation of large language models, 2021】能够以更高效、更经济的方式微调大型语言模型。NeMo-Aligner支持对多种对齐技术应用LoRA，我们遵循Wang等人【41, Helpsteer: Multi-attribute helpfulness dataset for steerlm, 2023】的训练流程，在Llama 2 70B模型上以及使用HelpSteer【41】和Open Assistant数据集【18】进行了SteerLM的LoRA训练。具体来说，我们将LoRA应用于所有注意力层，秩为32。我们使用全局批量大小128，AdamW优化器，经过10个预热步骤后使用恒定学习率1e-5，并训练了3个epoch。

效率与性能。如表6所示，将LoRA应用于BF16精度的SteerLM训练，可以将所需的最少80GB GPU数量从32个减少到8个。在相同GPU数量下，LoRA相比全参数微调实现了5倍的加速，同时保持了相当的模型性能：MT-Bench分数为7.43 vs. 7.54，这在该基准的噪声范围内【16, Mistral 7b, 2023】。

表6：遵循Wang等人 (2023) 训练流程的全参数和LoRA SteerLM比较。

LoRA扩展性。随着用于LoRA训练的GPU数量增加，相对吞吐量（以每秒样本数衡量）几乎成比例地提高，如图4所示。这表明NeMo-Aligner能够有效地在大量GPU上分配和并行化工作负载，且开销和收益递减效应极小。

6 SPIN训练

用SPIN复现SFT模型。我们使用SPIN【6, Self-play finetuning converts weak language models to strong language models, 2024】而非SFT来重新创建Zephyr-7B-Beta【39, Zephyr: Direct distillation of lm alignment, 2023】的SFT模型。我们从Mistral-7B基础模型【16, Mistral 7b, 2023】开始，并遵循Chen等人【6】的方法进行SPIN训练。然而，我们对他们的方法做了一些调整：我们没有将前一轮迭代的生成结果注入到当前迭代中（这会使数据集大小每轮翻倍），并且我们只训练一个迭代，每个迭代1个epoch。此外，我们使用了Ultrachat200k【10】数据集中一个仅包含5万样本的随机子集，而不是整个数据集，并使用了AdamW代替RMSProp。

训练配置与结果。我们的学习率为5e-7，总共400个步骤，其中40个为预热步骤，然后在最后100步使用余弦退火将学习率衰减到1e-7。全局批量大小为64，权重衰减为0.0，KL正则化系数为0.1，与Chen等人【6】一致。通过这种方法，我们取得了7.04的MT-Bench分数，超过了使用SFT的Zephyr-7B-Beta的6.64分【39】，也超过了3轮迭代的SPIN模型的6.78分【6】。

7 框架的可扩展性

可扩展设计。我们设计NeMo-Aligner时考虑了可扩展性，允许用户在分布式训练的复杂性下轻松修改算法。我们通过使用trainer抽象来实现这一点，该抽象鼓励在各种步骤和方法中重用现有的trainer方法。NeMo-Aligner的可扩展性使得DPO的变体，如Identity Preference Optimization (IPO)【1, A general theoretical paradigm to understand learning ´ from human preferences, 2023】、Conservative DPO (CDPO)【22, PyTriton: Framework facilitating NVIDIA Triton inference server usage in Python environments, 2022】和Kahneman-Tversky Optimization (KTO)【13, human-centerred loss functions (halos), 2023】，能够以最少的代码更改进行集成。此外，其他模型对齐技术，如Constitutional AI【3, Constitutional ai: Harmlessness from ai feedback, 2022b】、Rejection Sampling【37, Llama 2: Open foundation and fine-tuned chat models, 2023】和Self-Rewarding Language Models【46, Self-rewarding language models, 2024】，也正在通过框架设计被整合到NeMo-Aligner中。

A4 实验环境

模型架构:
- Llama 3 70B & 405B
- Llama 2 70B
- Nemotron 4 340B
- Mistral-7B
数据集:
- HelpSteer2: 用于RLHF PPO训练的提示来源。
- Ultrachat / Ultrafeedback: 用于DPO训练（复现Zephyr-7B-Beta），Ultrachat用于SFT，Ultrafeedback用于DPO。
- HelpSteer / Open Assistant: 用于SteerLM LoRA训练。
- MT-Bench: 用于评估所有对齐后模型的性能，使用GPT-4-Turbo作为裁判。
硬件配置:
- 计算节点: 每个节点包含8块H100-80GB-SXM GPU。
- 互联: 节点内使用NVLink，节点间使用Infiniband。
- 规模: 实验规模从12个节点（8 actor + 4 critic）扩展到126个节点（1008块H100 GPU），用于训练405B模型。
软件配置:
- 核心框架: NeMo-Aligner。
- 底层库: 基于Megatron-LM和NeMo构建。
- 优化工具: 集成了Transformer Engine的优化核、TensorRT-LLM用于推理加速、PyTriton用于分布式模型通信。
- 精度: 实验主要使用BF16精度。
- 优化器: 分布式融合Adam优化器，AdamW。

A4 实验结果

RLHF (PPO) 训练质量:
- 实验内容: 对Llama 3 70B SFT模型进行PPO训练。
- 实验结果: 在MT-Bench上，模型分数从SFT检查点的7.96提升到8.13。
- 结论: NeMo-Aligner能有效执行RLHF流程并提升模型质量。
RLHF (PPO) 可扩展性:
- 实验内容: 在不同节点规模下（从8+4到64+32节点）和不同负载下（全局批大小128 vs 1024）测试训练速度。
- 实验结果:
  - 从8+4节点扩展到16+8节点，整体速度提升1.80倍（见表1）。
  - 生成（Rollout）阶段扩展性近乎线性（2.05倍），训练（Train）阶段次线性。
  - 在1024的更大批大小下，系统能有效扩展至768块H100 GPU（64+32节点）（见表2）。
- 结论: NeMo-Aligner的分布式架构和优化能有效利用大规模计算资源，实现高效扩展。
RLHF (PPO) 系统性能消融研究:
- 实验内容: 逐一移除关键优化组件，测量对整体训练速度的影响。
- 实验结果: TensorRT-LLM集成是性能提升最关键的因素，带来近7倍加速。模型并行策略重分片（resharding）、TensorRT Refit和异步通信也分别带来3.87倍、3.15倍和1.54倍的显著加速（见表4）。
- 结论: NeMo-Aligner的各项系统优化都对整体性能有重要贡献。
最大规模模型训练:
- 实验内容: 使用1008块H100 GPU对Llama 3.1 405B模型进行PPO训练。
- 实验结果: 成功运行了训练流程，但速度慢于70B模型，瓶颈在于需要流水线并行的生成阶段（见表5）。
- 结论: NeMo-Aligner具备训练当前最大规模开源LLM的能力。
DPO 训练:
- 实验内容: 复现Zephyr-7B-Beta的训练流程。
- 实验结果: 最终模型MT-Bench分数为7.60，优于原始论文的7.34。
- 结论: NeMo-Aligner的DPO实现能够复现并超越现有SOTA结果。
SteerLM + LoRA 训练:
- 实验内容: 对Llama 2 70B模型进行全参数微调和LoRA微调的对比。
- 实验结果:
  - LoRA将所需GPU数量从32个减至8个，并在相同硬件上实现5倍加速（见表6）。
  - 模型性能与全参数微调相当（MT-Bench 7.43 vs. 7.54）。
  - LoRA训练的吞吐量随GPU数量增加呈近线性增长（见图4）。
- 结论: NeMo-Aligner高效支持PEFT方法，能以更低成本实现高质量对齐，且扩展性良好。
SPIN 训练:
- 实验内容: 使用SPIN方法训练Mistral-7B模型，替代Zephyr-7B-Beta中的SFT步骤。
- 实验结果: 获得MT-Bench分数7.04，优于原始SFT模型（6.64）和原始SPIN论文中的模型（6.78）。
- 结论: NeMo-Aligner中的SPIN实现是有效的，并能取得具有竞争力的结果。

A5 结论

现代模型对齐技术，特别是基于强化学习的技术，在系统实现方面带来了复杂的优化挑战。我们创建并开源了NeMo-Aligner，旨在让AI研究人员和从业者能够通过可扩展的方式，高效地利用所有可用计算资源进行LLM对齐实验。我们的框架在用更多计算资源训练大模型时表现出持续良好的扩展性。由于这是我们的初次发布，我们预计未来版本的扩展性只会更好。此外，我们支持SFT、PPO、DPO、SteerLM，并为计算资源有限的场景提供了使用LoRA进行参数高效微调的方式。作为一个采用Apache 2.0许可证的开源代码库，NeMo-Aligner可以使对齐研究更加高效和易于普及。

A6 附录

本文档不包含附录部分。

NeMo-Aligner: Scalable Toolkit for Efficient Model Alignment