重要提示

您正在查看 NeMo 2.0 文档。此版本对 API 和新库 NeMo Run 进行了重大更改。我们目前正在将所有功能从 NeMo 1.0 移植到 2.0。有关先前版本或 2.0 中尚不可用的功能的文档，请参阅 NeMo 24.07 文档。

重要提示

在开始本教程之前，请务必查看 简介，以获取有关设置 NeMo-Aligner 环境的提示。

如果您遇到任何问题，请参阅 NeMo 的已知问题页面。该页面列举了已知问题，并在适当的地方提供了建议的解决方法。

完成本教程后，请参阅 评估文档，以获取有关评估已训练模型的提示。

通过 RLHF 进行模型对齐#

在本教程中，我们将使用 NeMo Framework 中的模型完成整个人工反馈强化学习 (RLHF) 流程。这些模型可以包括 LLaMa 或 Mistral，并且我们的脚本将在它们之间保持一致的功能。

RLHF 通常先进行监督微调 (SFT)。我们应该首先遵循先决条件指南和 SFT 指南。获得 SFT 模型后，我们将使用它来启动 RLHF 过程。我们将使用 PPO 算法在 Anthropic-HH-RLHF 数据集上进行强化学习。

RLHF 的数据处理#

我们准备了一个脚本，可用于将 Anthropic-HH 数据集处理为 JSONL 格式。在 download_and_process.py 脚本上运行以下命令，以处理 anthropic HH。

python download_and_process.py

运行此脚本后，您应该拥有文件 {train,test}_comparisons.jsonl 和 {train,test}_prompts.jsonl。比较文件用于奖励模型训练，而提示文件用于强化学习训练。

奖励模型训练#

奖励模型用于对响应的质量进行评分。它使用成对比较损失进行训练，这意味着它需要响应对数据集。在每对中，一个响应的排名高于另一个响应。一个构造良好的奖励模型对于近端策略优化 (PPO) 训练过程的成功至关重要。

数据预处理#

您还可以为您自己的数据用于奖励模型训练阶段。奖励模型数据集需要以下格式

{"text": prompt1 || good_response_1}
{"text": prompt1 || bad_response_1}
{"text": prompt2 || good_response_2}
{"text": prompt2 || bad_response_2}
...

|| 表示字符串连接，而 *prompt1* 和 *prompt2* 是不同的提示。请注意，对于相同的提示，数据集中的 *prompt || good_response* 必须在 *prompt || bad_response* 之前。

JSONL 文件示例可能如下所示

{"text": "User: When did Virgin Australia start operating?\nAssistant: 31 August 2000"}
{"text": "User: When did Virgin Australia start operating?\nAssistant: I refuse to answer this question."}
{"text": "User: What is 6*10?\nAssistant: 60"}
{"text": "User: What is 6*10?\nAssistant: 90"}
...

要启动奖励模型训练，您必须从预训练或 SFT 训练的模型开始。在本节中，我们将使用上一步中训练的 SFT 模型来训练奖励模型。

要在终端上直接运行奖励模型训练

GPFS="/path/to/nemo-aligner-repo"
TRAIN_DATA_PATH="/path/to/train_comparisons.jsonl"
VALID_DATA_PATH="/path/to/test_comparisons.jsonl"

python -u ${GPFS}/examples/nlp/gpt/train_reward_model.py \
   trainer.num_nodes=1 \
   trainer.devices=8 \
   ++model.micro_batch_size=1 \
   ++model.global_batch_size=512 \
   ++model.data.data_impl=jsonl \
   pretrained_checkpoint.restore_from_path=/path/to/megatron_gpt_sft.nemo \
   "model.data.data_prefix={train: [${TRAIN_DATA_PATH}], validation: [${VALID_DATA_PATH}], test: [${VALID_DATA_PATH}]}" \
   exp_manager.create_wandb_logger=False \
   exp_manager.wandb_logger_kwargs.project=rm_training \
   exp_manager.wandb_logger_kwargs.name=rm_training \
   exp_manager.explicit_log_dir=/results

要使用 Slurm 运行奖励模型训练，请使用以下脚本。以下脚本使用 4 个节点，但您可以将节点计数更改为更小的数字。

#!/bin/bash
#SBATCH -A <<ACCOUNT NAME>>
#SBATCH -p <<PARTITION NAME>>
#SBATCH -N 4
#SBATCH -t 4:00:00
#SBATCH -J <<JOB NAME>>
#SBATCH --ntasks-per-node=8
#SBATCH --exclusive
#SBATCH --overcommit

GPFS="/path/to/nemo-aligner-repo"
PRETRAINED_CHECKPOINT_NEMO_FILE="/path/to/megatron_gpt_sft.nemo"

TRAIN_DATA_PATH="/path/to/train_comparisons.jsonl"
VALID_DATA_PATH="/path/to/test_comparisons.jsonl"

PROJECT="<<WANDB PROJECT>>"

CONTAINER=<<<CONTAINER>>> # use the latest NeMo Training container, Aligner will work there

MOUNTS="--container-mounts=MOUNTS" # mounts

RESULTS_DIR="/path/to/result_dir"

OUTFILE="${RESULTS_DIR}/rm-%j_%t.out"
ERRFILE="${RESULTS_DIR}/rm-%j_%t.err"
mkdir -p ${RESULTS_DIR}

MOUNTS="--container-mounts=MOUNTS" # mounts

read -r -d '' cmd <<EOF
echo "*******STARTING********" \
&& echo "---------------" \
&& echo "Starting training" \
&& cd ${GPFS} \
&& export PYTHONPATH="${GPFS}:${PYTHONPATH}" \
&& export HYDRA_FULL_ERROR=1 \
&& python -u ${GPFS}/examples/nlp/gpt/train_reward_model.py \
   trainer.num_nodes=${SLURM_JOB_NUM_NODES} \
   trainer.devices=8 \
   pretrained_checkpoint.restore_from_path='${PRETRAINED_CHECKPOINT_NEMO_FILE}' \
   "++model.data.data_prefix={train: [${TRAIN_DATA_PATH}], validation: [${VALID_DATA_PATH}], test: [${VALID_DATA_PATH}]}" \
   ++model.micro_batch_size=1 \
   ++model.global_batch_size=512 \
   ++model.data.data_impl=jsonl \
   exp_manager.explicit_log_dir=${RESULTS_DIR} \
   exp_manager.create_wandb_logger=True \
   exp_manager.wandb_logger_kwargs.name=${NAME} \
   trainer.rm.save_interval=500 \
   trainer.rm.val_check_interval=100 \
   trainer.rm.limit_val_batches=100000 \
   exp_manager.wandb_logger_kwargs.project=${PROJECT}
EOF

srun --no-container-mount-home -o $OUTFILE -e $ERRFILE --container-image=$CONTAINER $MOUNTS bash -c "${cmd}"
set +x

注意

目前，示例训练脚本不会自动在提供的测试集上运行评估。这可能会在未来的版本中更改。

在奖励模型训练期间，预计验证准确率会随着训练的进行而提高。在使用 Slurm 提供的上述示例中，我们实现了 69.57% 的验证准确率。

完成训练后，NeMo-Aligner 将保存一个 megatron_gpt.nemo 文件，该文件充当 RL 阶段所需的奖励模型。

PPO 训练#

在使用 SFT 微调 GPT 模型并训练奖励模型（如前节所述）后，您可以继续使用 PPO 进行 RLHF。

在 PPO 训练期间，我们概念上有四个模型相互交互

  1. PPO Actor 网络（也称为策略网络）：这是我们正在训练的模型。它应该从 SFT 模型开始。

  2. 奖励模型 (RM) 网络（也称为偏好模型）：此模型将提示与响应连接起来作为输入，并输出一个标量值，即奖励，PPO 算法将尝试最大化该奖励。

  3. PPO Critic 网络（也称为价值网络）：由于 PPO 是一种 Actor-Critic 算法，我们需要一个 Critic 来指导训练期间的 Actor。Critic 将为 Actor 提供的响应中的每个 token 提供价值估计。这些价值可以看作是 Actor 在生成所有剩余 token 后将收到的总奖励的估计值。Critic 应该从 RM 初始化，以便在训练的早期阶段提供有用的反馈。注意：RM 为整个序列生成单个奖励，而 Critic 为每个 token 生成一个值。

  4. 初始策略网络（也称为参考模型）：我们使用此模型来计算 KL 散度惩罚项，以确保 PPO Actor 不会与初始策略发散太多。这样，我们可以防止 PPO Actor 过拟合 RM 给出的奖励，并确保它不会忘记在预训练和 SFT 期间获得的知识。此模型应该是用于初始化 PPO Actor 网络的模型。

在最优配置中，NeMo-Aligner 将在同一作业中运行 Actor 和初始策略，并在同一作业中运行 Critic 和奖励模型。然后，它将使用 CPU 卸载在需要时加载回相应的模型。

下一节讨论如何启动这两个作业中的每一个。

启动奖励模型和 Critic 服务器#

要启动服务器

#!/bin/bash
# Example: If using the reward model trained from the above configuration, you can find
           the trained reward model checkpoint here: "/results/checkpoints/megatron_gpt.nemo"
CHECKPOINT_NEMO_FILE="/path/to/trained_rm.nemo"
GPFS="/path/to/nemo-aligner-repo"

RESULTS_DIR="critic_results_dir"

export PYTHONPATH="${GPFS}:${PYTHONPATH}" \
&& export HYDRA_FULL_ERROR=1 \
&& python -u ${GPFS}/examples/nlp/gpt/serve_ppo_critic.py \
   trainer.devices=8 \
   trainer.num_nodes=1 \
   ++model.tensor_model_parallel_size=1 \
   ++model.pipeline_model_parallel_size=1 \
   exp_manager.create_wandb_logger=False \
   exp_manager.wandb_logger_kwargs.name=critic_training \
   exp_manager.wandb_logger_kwargs.project=nemo_aligner_ppo \
   exp_manager.explicit_log_dir=${RESULTS_DIR} \
   trainer.ppo.inference_micro_batch_size=4 \
   ++pretrained_checkpoint.restore_from_path=${CHECKPOINT_NEMO_FILE} \
   ++model.megatron_amp_O2=True \
   ++model.activations_checkpoint_granularity=null \
   ++trainer.ppo.combine_rm_and_critic_server=True \
   ++model.offload_adam_states=True \
   ++model.mcore_gpt=True

上面的示例在 8 个 GPU 和 1 个节点上启动奖励模型 Critic 服务器。请务必根据您的模型大小和规模更改 trainer.devices、trainer.num_nodes。NeMo-Aligner 可在任何规模下工作。此外，请务必调整 trainer.ppo.inference_micro_batch_size 参数，因为这决定了 PPO Actor 允许发送到每个 DP 秩的 Critic 的批大小。

启动初始策略和 PPO Actor 训练#

PPO Actor 训练作业包含主控制器，该控制器在需要时对所有服务器进行 HTTP 调用。要启动 PPO Actor 和初始策略服务器

GPFS="/path/to/nemo-aligner-repo"
TRAIN_DATA_PATH="/path/to/train_prompts.jsonl"
VALID_DATA_PATH="/path/to/test_prompts.jsonl"

PRETRAINED_ACTOR_NEMO_FILE="/path/to/sft_checkpoint.nemo"

RESULTS_DIR="actor_results_dir"

export PYTHONPATH="${GPFS}:${PYTHONPATH}" \
&& export HYDRA_FULL_ERROR=1 \
&& python -u ${GPFS}/examples/nlp/gpt/train_gpt_ppo_actor.py \
   "++model.data.data_prefix={train: [${TRAIN_DATA_PATH}], validation: [${VALID_DATA_PATH}], test: [${VALID_DATA_PATH}]}" \
   ++model.data.data_impl=jsonl \
   pretrained_checkpoint.restore_from_path=${PRETRAINED_ACTOR_NEMO_FILE} \
   trainer.num_nodes=1 \
   trainer.devices=8 \
   ++model.pipeline_model_parallel_size=1 \
   ++model.tensor_model_parallel_size=1 \
   ++model.ppo.combine_rm_and_critic_server=True \
   ++model.ppo.offload_adam_states=True \
   ++model.megatron_amp_O2=True \
   ++trainer.ppo.normalize_advantages=True \
   ++model.mcore_gpt=True \
   exp_manager.create_wandb_logger=False \
   exp_manager.wandb_logger_kwargs.name=ppo_actor_training \
   exp_manager.wandb_logger_kwargs.project=nemo_aligner_ppo \
   exp_manager.explicit_log_dir=/rlhf/actor_test \
   ++model.ppo.entropy_bonus=0.0 \
   remote_critic_rm.pad_to_length=2048

上面的脚本在 1 个节点和 8 个 GPU 上启动初始服务器和 Actor 服务器。

注意

有关 PPO 超参数的更多信息，请参阅 PPO Hparams。

启动 RLHF 训练的两个服务器#

您可以使用 Slurm 启动两个作业，并使用以下脚本在完整的 RLHF 作业中协调它们

#!/bin/bash
#SBATCH -N 1 --ntasks-per-node 8 -A <<ACCOUNT>> -p <<PARTITION>> --job-name <<JOBNAME>> -t 4:00:00 --exclusive
#SBATCH hetjob
#SBATCH -N 1 --ntasks-per-node 8 -A <<ACCOUNT>> -p <<PARTITION>> --job-name <<JOBNAME>> -t 4:00:00 --exclusive

# To ensure determinism when calculating log probabilities between two forward-passes with identical weights, it is strongly
# recommended to set NCCL_ALGO. See https://github.com/NVIDIA/Megatron-LM/blob/b3375a0e38c10e2300ef4be031f7dcabab52b448/megatron/training/arguments.py#L593-L595
# for options.
export NCCL_ALGO=Tree

NAME="2p_ppo"

# PARAMETERS
RM_NEMO_FILE="/path/to/trained_rm.nemo"

ACTOR_NEMO_FILE="/path/to/sft_model.nemo"

TRAIN_DATA_PATH="/path/to/train_prompts.jsonl"
VALID_DATA_PATH="/path/to/test_prompts.jsonl"

RESULTS_DIR="/path/to/results_dir"
mkdir -p $RESULTS_DIR

GPFS="/path/to/nemo-aligner-repo"
MOUNTS="--container-mounts=MOUNTS" # mounts

CONTAINER=<<<CONTAINER>>> # use the latest NeMo Training container, Aligner will work there

PROJECT=ppo_run

CRITIC_LOG_DIR="${RESULTS_DIR}/critic_results"
CRITIC_OUTFILE="${CRITIC_LOG_DIR}/critic_output_%j_%t.log"
CRITIC_ERRFILE="${CRITIC_LOG_DIR}/critic_error_%j_%t.err"
CRITIC_PORT=5567

mkdir -p $CRITIC_LOG_DIR

CRITIC_NAME="${NAME}_critic"

read -r -d '' cmd_critic_inference <<EOF
cd ${GPFS} \
&& export PYTHONPATH="${GPFS}:${PYTHONPATH}" \
&& export HYDRA_FULL_ERROR=1 \
&& python -u ${GPFS}/examples/nlp/gpt/serve_ppo_critic.py \
   trainer.ppo.inference_micro_batch_size=4 \
   trainer.devices=8 \
   trainer.num_nodes=${SLURM_JOB_NUM_NODES_HET_GROUP_0} \
   exp_manager.explicit_log_dir=${CRITIC_LOG_DIR} \
   exp_manager.create_wandb_logger=True \
   exp_manager.wandb_logger_kwargs.name=${CRITIC_NAME} \
   exp_manager.wandb_logger_kwargs.project=${PROJECT} \
   trainer.ppo.port=${CRITIC_PORT} \
   ++model.offload_adam_states=True \
   ++model.micro_batch_size=1 \
   ++model.global_batch_size=64 \
   pretrained_checkpoint.restore_from_path=${RM_NEMO_FILE}
EOF

srun --no-container-mount-home --het-group=0 -o $CRITIC_OUTFILE -e $CRITIC_ERRFILE --container-image=${CONTAINER} $MOUNTS bash -c "${cmd_critic_inference}" &

sleep 30

ACTOR_LOG_DIR="${RESULTS_DIR}/actor_results"
CHECKPOINT_DIR="${ACTOR_LOG_DIR}/checkpoints"
TENSOBOARD_DIR="${ACTOR_LOG_DIR}/tensorboard"

PPO_ERRFILE="${ACTOR_LOG_DIR}/actor_error_%j_%t.err"
PPO_OUTFILE="${ACTOR_LOG_DIR}/actor_output_%j_%t.log"

mkdir -p $ACTOR_LOG_DIR
mkdir -p $TENSOBOARD_DIR
mkdir -p $CHECKPOINT_DIR

ACTOR_NAME="${NAME}_actor"

host_critic="$(scontrol show hostnames=$SLURM_JOB_NODELIST_HET_GROUP_0 | head -n1)"

read -r -d '' cmd_ppo <<EOF
cd ${GPFS} \
&& export PYTHONPATH="${GPFS}:${PYTHONPATH}" \
&& export HYDRA_FULL_ERROR=1 \
&& python -u ${GPFS}/examples/nlp/gpt/train_gpt_ppo_actor.py \
   trainer.devices=8 \
   trainer.num_nodes=${SLURM_JOB_NUM_NODES_HET_GROUP_1} \
   trainer.ppo.max_steps=15 \
   ++model.data.data_impl=jsonl \
   "++model.data.data_prefix={train: [${TRAIN_DATA_PATH}], validation: [${VALID_DATA_PATH}], test: [${VALID_DATA_PATH}]}" \
   pretrained_checkpoint.restore_from_path=${ACTOR_NEMO_FILE} \
   exp_manager.explicit_log_dir=${ACTOR_LOG_DIR} \
   exp_manager.create_wandb_logger=True \
   exp_manager.wandb_logger_kwargs.name=${ACTOR_NAME} \
   exp_manager.wandb_logger_kwargs.project=${PROJECT} \
   ++model.micro_batch_size=1 \
   ++model.global_batch_size=64 \
   ++model.activations_checkpoint_granularity=selective \
   ++model.activations_checkpoint_method=uniform \
   ++model.optim.lr=9e-7 \
   trainer.ppo.val_check_interval=3 \
   ++model.optim.sched.min_lr=9e-8 \
   ++model.ppo.entropy_bonus=0.0 \
   ++model.ppo.ratio_eps=0.2 \
   ++model.ppo.num_rollout_samples=512 \
   ++model.ppo.rollout_micro_batch_size=8 \
   ++model.ppo.length_params.max_length=1024 \
   trainer.ppo.initial_policy_kl_penalty=0.02 \
   remote_critic_rm.critic.ip=${host_critic} \
   remote_critic_rm.critic.port=${CRITIC_PORT}
EOF

srun --no-container-mount-home --het-group=1 -o $PPO_OUTFILE -e $PPO_ERRFILE --container-image=${CONTAINER} $MOUNTS bash -c "${cmd_ppo}" &

wait

上面的脚本在 1 个节点上运行奖励模型 Critic 服务器，在 1 个节点上运行 Actor。

务必在 srun 命令后使用 & 启动所有作业，以确保它们不会相互阻塞。

注意

确保更改 Critic 参数 trainer.ppo.inference_micro_batch_size，以使 trainer.ppo.inference_micro_batch_size * DP size <= model.ppo.rollout_micro_batch_size。

使用 TensorRT-LLM 加速 PPO#

NeMo-Aligner 支持使用 TensorRT-LLM 加速 RLHF。启用后，这可以显着加速 PPO 训练时间。将 TensorRT-LLM 与 Aligner 结合使用时，需要设置几个关键标志。

  1. trainer.ppo.trt_llm.enable=True 启用 TensorRT-LLM。

  2. trainer.ppo.trt_llm.reshard=True 在推理期间启用动态流水线并行重新分片。这实际上是在流水线并行下运行训练，但在张量并行下仅进行推理。

  3. trainer.ppo.trt_llm.unload_engine_train=False：可以设置为 True 以在训练期间卸载引擎，这将释放更多内存用于训练，因为引擎位于 GPU 上，但这会增加加载和卸载引擎的成本。建议将其保持为 False。

  4. trainer.trt_llm.model_type=llama 告诉 TensorRT-LLM 我们要运行推理的模型类型。必须为其他模型类型更改，例如，对于 nemotron，这应为 gptnext。

  5. trainer.ppo.batch_iterator.use_flask=True 启用 flask 服务器以平衡不同 DP 工作程序之间的工作负载。

有关更多信息，请参阅 NeMo-Aligner 论文。

注意

如果您正在使用 TensorRT-LLM 加速交互式运行 train_gpt_ppo_actor.py（在 SLURM 外部），则必须在 python 运行命令前添加 mpirun -n 8 --allow-run-as-root

mpirun -n 8 --allow-run-as-root python -u ${GPFS}/examples/nlp/gpt/train_gpt_ppo_actor.py ...

如果您使用的是 SLURM，则无需添加 mpirun，因为如果您使用 --mpi=pmix 运行 srun，则会自动处理此问题

read -r -d '' cmd_ppo <<EOF
cd ${GPFS} \
&& export PYTHONPATH="${GPFS}:${PYTHONPATH}" \
&& export HYDRA_FULL_ERROR=1 \
&& python -u ${GPFS}/examples/nlp/gpt/train_gpt_ppo_actor.py \
    ...
EOF

srun --mpi=pmix ... bash -c "${cmd_ppo}"

使用 TensorRT-LLM 的 PPO 结果#

我们使用 TensorRT-LLM PPO 集成来训练 LLaMa3-70B-PPO-Chat 模型。此模型使用全局批大小 128、num_rollout_samples 128、恒定学习率 1e-7 和 KL 惩罚 3e-3 进行训练。有关更多详细信息，请参阅 Helpsteer2 论文。

使用 TensorRT-LLM 的 PPO 性能结果#

我们通过在 Helpsteer2 提示上运行 Llama3 70B Actor 和 Llama3 70B Reward 模型来测试我们的 TRT-LLM 集成的扩展性，其中 num_rollout_samples=1024、global_batch_size=128，启用重新分片，并将引擎卸载设置为 False。

Actor 节点计数

Critic 节点计数

生成的 Token 数量

步长时间（秒）

扩展效率

8

4

350.7

366.7

1

16

8

323.3

190.4

1.93

32

16

331.7

108

3.40

64

32

334

56.9

6.44

注意

对于 64x32 配置，由于分布式优化器的额外内存，我们使用了 16 的 rollout_micro_batch_size 而不是 8。

我们还支持在 Llama3.1 405B Actor 和 Reward 模型上运行 RLHF。以下数字是在 num_rollout_samples=128、global_batch_size=128、关闭重新分片、将引擎卸载设置为 False 的情况下生成的。

Actor 节点计数

Critic 节点计数

生成的 Token 数量

步长时间（秒）

84

42

915.6

164.6

未来，我们的目标是提高具有高流水线并行大小的大型模型的生成性能。

PPO 结果#

完成 RLHF 训练后，您可以使用 NeMo 代码库中的 megatron_gpt_eval.py 脚本来运行对已训练模型更严格的评估。

将本教程扩展到更大的模型#

虽然上面的教程提供了一种开始使用 RLHF 的方法，但它并不代表最优的性能或收敛配置。在完全运行 RLHF 时，我们预计 MT-bench 分数将提高约 +0.4 到 +0.5。从高质量的 SFT 模型开始并密切监控响应长度至关重要。

在本页中