NVIDIA Nemo Framework 高性能的强化学习训练

摘要：导读NVIDIA 在大语言模型训练领域的最新进展之一是提供了一个全面的技术栈——NeMo，用于支持大型语言模型（LLM）端到端的训练，包括强化学习（Reinforcement Learning）方面的支持。

导读NVIDIA 在大语言模型训练领域的最新进展之一是提供了一个全面的技术栈——NeMo，用于支持大型语言模型（LLM）端到端的训练，包括强化学习（Reinforcement Learning）方面的支持。

本次分享的主要内容包括：

1. NVIDIA 全栈：端到端大型语言模型（LLM）软件解决方案

2. 模型对齐与强化学习的重要性

3. NeMo Aligner 概要介绍

4. 基于人类反馈的强化学习（RLHF）概括

5. PPO 算法讲解

6. 强化学习（RLHF）在实现高性能方面面临的挑战

7. NeMo Aligner 核心特性

8. 优化效果

9. 未来工作与大模型训练的创新方向

10. Q&A

分享嘉宾｜高文雯 NVIDIA Lead Product Manager

编辑整理｜杨峰

内容校对｜李瑶

出品社区｜DataFun

首先来整体介绍一下 NVIDIA 技术栈。

上图中 x 轴代表了大语言模型训练的各个阶段。其中，最初的阶段是预训练（Pretrain），随后进入微调（Fine-tuning）阶段。在微调阶段，可以采用 PEFT 参数高效微调的方法，或者进行全量微调 Full FT 以及模型对齐。这些步骤旨在提升模型在特定任务或领域上的性能。最后进入 Inference 阶段，这是模型将学到的知识应用于实际任务的过程。

y 轴表示的是整体技术栈的层次结构，从底层到高层。在训练和微调方面，最底层是 Transformer Engine，这是开源的一个库。它提供了高性能的内核优化，并支持 FP8（8 位浮点数）格式，以便在先进架构上实现更高性能的训练。

在 Transformer Engine 之上，我们构建了 Megatron-core。这里将 Megatron-LM 中的 GPU 优化、并行策略等核心优势整合到了一个名为 Megatron-Core 的库中，从而进一步提升了训练和推理的效率。

在此之上就是 NeMo 框架，其涵盖了多个领域，包括大语言模型、多模态模型（无论是视觉还是视频），以及语音识别技术（如 ASR 和 TTS）。这些领域的融合使得 NeMo 框架在多个应用场景中都具有强大的竞争力。

在 NeMo 框架中，有几个重要的组件，其中一个名为 NeMo Aligner，是今天介绍的重点。NeMo Aligner 包含了针对模型对齐 RLHF 等算法的高性能的优化。这些优化旨在提升模型在复杂任务上的性能和稳定性。

技术全栈组成：

预训练（Pretrain）& 微调（Fine-tune）

NeMo Framework：NVIDIA 的大型语言模型预训练框架，用于训练大规模的语言模型, 以及微调预训练模型。这包括参数高效微调（PEFT）和全参数微调（Full FT），以及模型对齐。NeMo Aligner：NeMo 框架中的一个组件，专门用于模型对齐，确保模型输出与人类偏好一致。Megatron-Core：核心库，为训练提供基础支持，提高 GPU 的利用率。Transformer Engine/FP8：底层库，高性能算子，尤其支持 FP8 精度计算。

推理（Inference）

TRT & TRT-LLM：TensorRT（TRT）和 TensorRT-LLM 是用于推理阶段的优化工具，它们与 Triton 后端集成，以实现快速且高效的模型部署。

模型对齐和强化学习是提升大型语言模型效果的关键技术。通过模型对齐，我们可以确保模型的输出与人类的价值观和偏好相一致。而强化学习，尤其是基于人类反馈的强化学习，为模型提供了一种学习和适应人类偏好的有效机制。近期 GPT-o1 以及国内多家大型科技公司相继发布了关于推理及逻辑方面的模型。这些模型均巧妙地运用了强化学习技术，以进一步提升其逻辑能力。强化学习本身也因此变得尤为重要，它不仅在模型对齐方面发挥着关键作用，而且在提升推理能力方面也具有重要意义。

NeMo Aligner 是一个强大的工具，它通过集成多种最先进技术，提供了高性能和高吞吐量的模型微调能力。其设计考虑了模型的可扩展性和商业应用的可行性，同时提供了丰富的资源来支持研究和开发。发布这些算法的目的是希望它们能够在 GPU 上发挥出最佳性能，从而帮助用户充分利用 GPU 资源。为此，我们进行了大量的性能优化工作，包括模型并行、推理加速、FP8 长序列支持以及 sequence packing 等策略。这些策略的实施旨在提升算法在 GPU 上的运行效率，从而为用户提供更好的使用体验。

1. 最先进技术（SOTA Techniques）

NeMo Aligner 采用了一系列最先进技术，包括：

SFT（Supervised Fine-Tuning）：监督式微调，用于提高模型在特定任务上的性能。SteerLM（Conditioned SFT）：条件化监督微调，进一步优化模型以满足特定条件或指令。RLHF（Reinforcement Learning from Human Feedback）：基于人类反馈的强化学习，用于训练模型以更好地符合人类偏好。DPO | IPO | RPO：不同的优化策略，用于提高模型的决策过程。Constitutional AI：遵循宪法原则的人工智能，确保模型的输出符合特定的道德和法律标准。SPIN：一种用于提高模型解释性和透明度的技术。

2. 性能与可扩展性（Performance & Scalability）

NeMo Aligner 通过一系列优化提高性能和可扩展性，主要包括：

模型并行（Model Parallelism）：允许模型在多个处理器上并行处理，提高训练速度。加速推理（Accelerated Inference）：优化模型以加快推理过程，减少延迟。FP8：使用8位浮点数进行计算，减少内存使用并提高计算效率。长序列支持（Long Sequence Support）：处理和生成更长的序列数据。序列打包（Sequence Packing）：一种技术，用于更有效地将不同长度的微调训练数据打包成一个序列来提高训练效率。

3. 贡献开源社区

NVIDIA 积极贡献于开源社区，通过模型 checkpoints 和数据集（datasets），鼓励研究人员和开发者使用和尝试这些资源。这不仅促进了知识的传播，也为社区的创新和发展提供了基础。

04基于人类反馈的强化学习（RLHF）

在 NeMo Aligner 中，基于人类反馈的强化学习（RLHF）被认为是最复杂的算法。RLHF 的训练过程遵循 Ouyang 等人（2022 年）的研究，旨在通过人类反馈来优化模型的行为和输出。本文将详细阐述 RLHF 的三个主要训练阶段：监督式微调、奖励模型训练和近端策略优化。

1. 监督式微调（Supervised Fine-Tuning）

第一阶段是传统的 SFT（监督微调，Supervised Fine-Tuning）过程。在这个过程中，我们采用基础模型，并结合自定义数据，通过监督学习的方法对其进行微调，从而得到一个 SFT 模型。这一步骤相对直观，也是当前机器学习领域较为常见的做法。

预训练模型：使用可训练的权重开始。注释数据：与特定任务相关的数据集。SFT 模型：通过在注释数据上进行微调，生成一个 SFT 模型。数据集：用于微调的数据集。

2. 奖励模型训练（Reward Model Training）

第二阶段需要训练一个奖励模型（reward Model）。奖励模型的作用是评估模型生成的回答质量。具体而言，我们利用第一步得到的 SFT 模型生成一些回答，并将这些回答与对应的 prompt 组成数据集，即偏好数据集（Preference Dataset）。在这个数据集中，每个 prompt 下都有两个回答，人类标注者会根据这两个回答的质量，选择出较好的一个，并为其分配更高的奖励分数（reward），而另一个则获得较低的分数。通过这种方式，我们可以训练出一个能够区分回答质量的奖励模型。这一阶段包括：

SFT 模型：作为奖励模型训练的基础。偏好数据：包含人类偏好的示例数据。奖励模型：通过学习偏好数据来评估模型输出的优劣。

3. 近端策略优化（Proximal Policy Optimization）

第三阶段是是将最初的模型、经过 SFT 微调后的模型以及刚刚训练出的奖励模型等所有组件融合在一起，通过微调模型生成回复，然后奖励模型给予奖励，以此来进一步优化微调模型，并且同时不让其和最初使模型偏离过多。如此反复，以得到一个我们期望的最终模型。这一步相对复杂，因为它涉及到多个模型的整合和优化。这一阶段包括：

SFT模型：作为策略网络的起点。提示（Prompts）：用于生成模型响应的输入。策略网络（Actor）：初始于 SFT 模型，生成模型的输出，然后通过得到的奖励进一步训练。价值网络（Critic）：初始于奖励模型，评估策略网络输出的预期奖励。训练策略模型：通过结合奖励模型的反馈，优化策略网络以提高输出质量。

总之：RLHF 是一个多阶段的训练过程，它通过监督式微调、奖励模型训练和近端策略优化来提升模型的性能。这种方法使得模型能够更好地理解人类反馈，并据此优化其行为。通过这种复杂的训练过程，NeMo Aligner 能够训练出更符合人类偏好和价值观的模型。

近端策略优化（PPO）算法是一种高效的强化学习策略，它通过平衡策略更新的幅度和奖励最大化来提高学习稳定性和效率。以下将详细阐述 PPO 算法的关键步骤和组件。

PPO 算法的训练过程涉及以下几个核心步骤：

首先，我们会有一个特定的 prompt，这个 prompt 会同时输入给两个模型：一个是 SFT（监督微调）模型，另一个是基于 SFT 模型构建的 actor 模型。这两个模型在训练初期是完全相同的，因此它们会对相同的 prompt 产生回复（response）。随后，这些回复会被输入给 reward 模型和 critic 模型。这两个模型在训练初期也是基于我们之前训练好的 reward 模型构建的，因此它们会对回复进行评分。这个评分过程实际上是一个反馈机制，用于指导 actor 模型的训练。

具体来说，如果 actor 模型的回复得到了较高的评分，就说明这个回复是符合我们期望的，因此我们会给予这个模型一个正向的反馈，鼓励它继续生成类似的回复。相反，如果回复的评分较低，那么说明这个回复不符合我们的期望，因此我们会给予模型一个负向的反馈，让它调整自己的生成策略。

此外，为了确保 actor 模型在训练过程中不会偏离我们的初衷，我们还会引入一个反馈机制。这个机制会将 actor 模型与最初的模型进行比较，确保 actor 模型在训练过程中不会过度偏离其初衷。除了上述反馈机制外，还有一个重要的反馈循环存在于 Critic 模型和 reward 模型之间。我们希望 Critic 模型能够与 reward 模型保持对齐，因此它们之间也会有一个反馈循环。通过不断地调整 Critic 模型的参数，我们可以使其与 reward 模型的评分更加一致。

在训练过程中，我们会从样本中抽取 n 个 prompt，并生成对应的回复。然后，我们会计算每个回复的 reward 值以及 Critic 模型的评分。接着，我们会根据这些反馈对 actor 模型和 Critic 模型进行更新。这个过程会不断地循环进行，直到我们得到满意的模型性能。

值得注意的是，我们的系统中一共有四个模型：两个用于训练和推理的模型（即 actor 模型和 Critic 模型），以及两个只用于推理的模型（即最初的模型和 reward 模型）。其中，actor 模型和最初的模型具有相同的架构，而 Critic 模型和 reward 模型也具有相同的架构。在训练过程中，这些模型的参数会逐渐演变，以适应我们的训练目标。

大模型训练是一个复杂而精细的过程，需要我们在多个模型之间建立有效的反馈机制，并不断地调整和优化模型的参数。通过利用 GPU 等高性能计算资源，我们可以更加高效地实现这一目标，并得到我们期望的模型性能。在接下来的内容中，我们将进一步探讨如何设计有效的架构来最大限度地发挥 GPU 的效果。

1. RLHF 的复杂性

RLHF 的复杂性主要体现在以下几个方面：

需要协调四种不同的模型：policy 模型(即 Actor 模型)vs initial policy 模型（即微调模型）、critic 模型 vs reward 模型。policy 模型和 critic 模型都需要执行训练和推理任务。每个模型具有完全不同的性能特征。

2. NeMo Aligner 的主要目标

NeMo Aligner 的主要目标之一是最大化性能。据我们所知，NeMo Aligner 的 RLHF 实现速度超过了其他任何开源 RLHF 实现。

3. NeMo Aligner 的性能优化

NeMo Aligner 通过以下关键性能优化措施，实现了快速的 RLHF：

分布式架构：确保在模型的复杂协调中，所有 GPU 都能得到充分利用。这种架构允许多个 GPU 并行工作，提高了计算效率。TRT-LLM 快速推理：利用 TensorRT-LLM（TRT-LLM）进行快速推理，减少了模型在执行任务时的延迟。动态负载平衡：确保每个 GPU 都能接收到足够的工作量，以达到饱和状态，并且工作量在 GPU 之间平衡，以实现良好的扩展性。

接下来将详细介绍这三方面的优化。

首先聚焦于分布式系统架构的设计。具体而言，我们面对的是四个模型，它们被巧妙地分为两对：一对是 critic model 和 reward model，另一对则是初始的 actor model 和基准模型。值得注意的是，尽管每一对内部的模型大小相同（仅参数不同），但两对之间的模型大小可能存在显著差异。例如，一对可能采用 70B 参数的模型，而另一对则可能采用 7B 或 13B 的模型。

针对这种模型大小的不均衡性，我们采取了灵活的 GPU 资源分配策略。具体而言，我们根据模型的大小来动态调整 GPU 的分配。较大的模型获得更多的 GPU 资源，而较小的模型则得到相对较少的资源。这种策略不仅提高了 GPU 资源的利用率，还确保了根据模型的实际需求进行灵活的资源调度。

进一步地，对于每一对模型，我们采用了共享 GPU 资源的机制。这意味着，在需要时，我们会将所需模型的参数（weights）加载到 GPU 中，而在不需要时，则将这些参数卸载到主机内存（host memory）中。这种“按需加载”的策略进一步提升了 GPU 的使用效率，确保了没有任何 GPU 资源处于空闲状态。因为每一对模型中的两个模型总有一个是处于工作状态，当轮到某个模型工作时，我们会迅速将其参数加载到 GPU 中，并利用分配给它的资源进行计算。

此外，在两对模型之间，我们采用了 PyTorch Triton 进行调度。这种调度机制使得两个任务（job）在大多数情况下都可以独立运行，仅在每个循环的最后阶段进行同步。为了协调这两个任务，我们设计了一个调度器（scheduler），以确保两对模型之间的协同工作。例如，actor model 会生成一些响应（response），这些响应会被发送给左边的模型对。当左边的模型对积累了足够多的响应后，会进行批量推理（batch inference），并为这些响应计算评分。在这个过程中，PyTorch Triton 作为服务器（server）处理右边的模型对，而左边的模型对则作为客户端（client）与之交互。

上图说明了在两组 cluster 上的运行方式。首先，policy model 和 initial model 负责生成响应；另一组是 Critic model 和 reward model，接收生成的响应并进行推理以给出评分。有时生成过程可能较长，Critic 和 reward 模型就需要短暂的等待。一旦生成过程结束，Critic 和 reward 模型便立即开始推理评分，与此同时，actor model 和 initial state model 也进入训练阶段。为了确保训练过程的高效性，我们根据两组模型的规模动态分配资源，以确保它们都能获得足够的计算资源，同时避免任何一方的等待。我们的目标是最大化资源利用率，使两组模型几乎能够同时完成训练。

然而，在实际运行过程中，偶尔会出现一方模型先于另一方完成训练的情况。例如，如果 actor model 和 initial state model 先完成训练，它们将立即进入下一个生成循环，而 Critic 和 reward 模型则继续训练，直至完成。在这种情况下，当 Critic 和 reward 模型训练完毕后，它们将使用 actor model 新生成的响应进行下一轮的评分。为了应对这种偶尔出现的不同步情况，我们设计了一种调节机制，以确保两组模型能够获得恰当的资源，从而使它们的训练过程尽可能同步。这种机制不仅减少了 GPU 的等待时间，还提高了整体训练效率。

生成所需的计算资源是较为昂贵的，并且会占据 60% 甚至 90% 以上的时间，因此我们希望利用 TensorRT-LLM 进行推理操作，以加速生成过程。

具体流程如下：

训练完成后，将模型复制并转化为推理模型。推理模型负责生成任务，生成的数据随后用于训练。训练模型接收新生成的数据，形成循环。

在训练和推理过程中，内存管理尤为关键。训练过程中涉及模型权重、激活值以及优化器状态等，而推理则需要加载推理引擎及其权重。为了优化内存使用，采取以下措施：

在推理前，将训练相关的激活和优化器状态卸载到主机内存。推理所需的 TRT 引擎加载到 GPU 中，开始生成过程。生成完成后，将TRT引擎的数据移回主机内存，并将训练所需的梯度和优化器状态调入 GPU 内存。

具体而言，在推理开始之前，我们会将训练相关的数据（如激活值和优化器状态）卸载到主机内存中，以释放显存空间供推理引擎使用。当推理完成后，我们再将推理引擎相关的数据卸载到内存中，并将训练所需的数据（如梯度和优化器状态）加载到显存中，以便开始下一轮的训练。这一过程需要精细的显存管理，以确保两个任务在有限的资源下高效运行。

此外，还注意到，对于较小的模型，推理引擎本身可能并不占用大量显存。在这种情况下，卸载操作可能并不划算，因为卸载本身也需要一定的时间和资源开销。因此，在实际应用中，我们需要根据模型的大小和显存资源来权衡是否进行卸载操作。

在卸载过程中，我们也应用了一些底层技术来优化性能，如 Pinned memory 和异步重叠（async overlap）等。这些技术有助于最大限度地利用显存和计算能力。

训练完成后，新生成的模型权重通过 TensorRT 的 refit 技术推送到推理侧。Refit 允许在不重建引擎的情况下，直接更新现有引擎的权重，显著减少了训练和推理的时间。

综上，我们在一套资源上实现了训练和推理的并行运行，并通过动态显存管理、底层技术优化以及 TensorRT refit 等技术手段，有效提高了大模型训练的效率。

推理过程相比训练过程使用更少的内存，因为不需要优化器状态或梯度，因此不需要像训练时那样高的模型并行性。

这里采用了 reshard 技术来降低模型的并行需求。重点是减少通信开销，例如，管道并行（Pipeline Parallelism, PP）可能会引入 30% 的开销，所以我们尽量减少甚至去掉管道并行。

另外，在推理过程中，卸载优化器的权重和梯度，从而减少内存占用。

最后，TRT-LLM 引擎可以在内存中保持，或者在训练期间卸载，以优化资源使用。

在大规模模型训练过程中，负载均衡是一个至关重要的考量因素。特别是当我们采用数据并行（Data Parallelism）策略时，如何有效分配和管理各个数据处理器（如 GPU）上的工作负载，成为了一个亟待解决的问题。

数据并行工作的问题：

数据并行工作节点处理提示（prompts）以填充一个推理的批次（rollout batch）。如果给每个 DP 工作节点分配相同数量的提示，可能会导致一些节点空闲。一些提示可能因为生成长度较长而需要更长时间，这会迫使其他数据并行工作节点等待。随着训练规模的扩大，不同 DP 机器（ranks）之间的工作负载差异将会增加。

动态负载平衡的解决方案：

使用 Flask server 在 DP 等级 0 上接收来自其他等级的 prompts。每个 DP 等级在完成生成后请求新的 prompts。通过在运行时动态调整 micro batch 的大小或每个数据并行实例的样本数量，确保工作负载始终平衡。另外，在 TRTLLM 的 inflight batching 中，也动态地在推理 micro batch 内加载 batch，以提高资源利用率。08优化效果

采用 TensorRT-LLM 进行推理加速，我们实现了接近 7 倍的速度提升。这一成果显著降低了模型 RLHF 的时间成本，为实际应用提供了强有力的支持。

此外，我们还采用了两种额外的优化手段：动态负载均衡和分布式系统。这两种技术的结合使用，进一步带来了大约两倍的性能提升。因此，相较于未经优化的原始训练代码，我们的总体性能提高了近 15 倍。以前需要一周时间的训练现在可以在一夜之间完成。这对于大规模强化学习任务尤为重要，可以显著提高研究和开发的效率。

上图以更直观的方式展示了优化技术的实际效果。y 轴代表时间，数值越小表示性能越好。

我们将结果分为四大类进行展示。以紫色区域为例，它代表如果没有 Refit 的作用，时间将显著增加。类似的，如果不做 Reshard，粉色区域，也就是在推理阶段时间会大幅增加，如果没有 TRT-LLM，就会增长更多，可能达到七八倍。如果不做 Async training，性能也会受到一定影响。

相同的模型部署更大规模的机器资源上，可以达到 1.8 倍的 scaling，这是非常可观的提升。

我们特别针对 LLaMA 405B 模型进行了实验。我们采用了 84 个节点作为 actor，42 个节点作为 critic，结果如上图所示。

最后我们可以得出结论，NeMo Aligner 可以达到极高的速度。近期，尽管可能有一些其他发表的论文提到使用 NeMo 进行大模型训练的性能测试，但可能其中并未采用最新版本的 NeMo，因此展现的并非最新数据。NeMo 本身也有报告论文发表，大家可以查阅并尝试复现其中的实验。

NeMo 的优势不仅限于速度方面，其拓展性同样出色。在支持上千个 GPU 的集群环境中，NeMo 能够稳定运行并发挥出卓越的性能。这得益于平台内部的多项优化技术，例如之前提到的并行技术，以及针对推理阶段的各种异步训练优化等。这些优化细节使得 NeMo 在处理大规模模型训练时更加高效。

以具体的例子来说，在使用 NeMo 训练 LLaMA 3 70B 做 RLHF，可以在两个小时内完成，远小于其他训练方案所需时间。这一显著差异充分展示了 NeMo 在大模型训练领域的强大实力。

09未来工作与大模型训练的创新方向

Nvidia 在大模型训练领域持续探索，并致力于推出更多创新技术和方法。在推理方面，我们正积极寻求更高效的方法，以加速 engine 从 GPU 到 CPU 的加载过程。同时，针对并行计算中普遍存在的较大开销问题，我们也在探索如何进一步提升并行计算的速度和效率。

在训练方面，我们关注到了 sequence packing 的重要性。由于 prompt 和 response 的长度可能差异较大，我们采用了多种策略来优化 sequence packing。一种方法是动态的负载均衡，即根据数据长度动态调整分配；另一种方法是在一个 batch 中尽量将长度相近的数据打包在一起，以减少计算开销。此外，针对当前长序列处理的趋势，我们利用 ring attention 等技术来做 context parallel，以更有效地处理长序列数据。在低精度训练方面，我们也进行了相关实验，旨在探索在保证模型性能的同时，进一步降低计算成本和能耗。

除了上述工作外，我们还进行了知识蒸馏的研究，并在模型对齐过程中引入了新的算法，以提升模型的准确性和效率。同时，我们还在探索跨模态模型的应用，以拓展应用场景和范围。

上图中列出了一些相关资源供大家参考，欢迎大家关注及进一步探讨。

10Q&A

Q：关于大模型训练，PyTorch 也在进行一系列工作，如 2.0 版本中的编译功能以及分布式的自动并行能力等，这些优化工作是否会对 NeMo 产生影响？

A：PyTorch 致力于提供一个相对通用的解决方案，因此其优化可能并不针对于某一特定领域。而 NeMo 框架在模型对齐方面进行了大量的优化，甚至针对 RLHF 本身也进行了优化，这也是为什么它能够实现近 15 倍的性能提升。

对于刚开始进行大模型训练实验的研究人员来说，使用 PyTorch 作为入门工具是没有问题的，尤其是当模型规模较小、效果尚需验证时。然而，随着实验的深入和模型复杂度的增加，我们鼓励大家尝试 NeMo 这一更专注于特定领域优化的框架。

同时，Nvidia 一直致力于为所有主流工具提供最优的 GPU 加速解决方案，所以 Nvidia 也有针对 PyTorch 加速的团队。我们希望将 NeMo 中的优化也逐渐落地到 PyTorch 中，以便大家在入门时就能获得出色的体验。因此，我们采取的是双管齐下的策略，既提供通用的解决方案，又针对特定领域进行优化。

以上就是本次分享的内容，谢谢大家。