重要提示

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

通过 REINFORCE 进行模型对齐#

在本教程中，我们将指导您完成使用 REINFORCE 对齐 NeMo 框架模型的过程。此方法可以应用于各种模型，包括 LLaMa2 和 Mistral，并且我们的脚本在不同模型之间保持一致的功能。

REINFORCE 通常先于监督微调 (SFT)。我们应该首先遵循先决条件指南和 SFT 指南。获得 SFT 模型后，我们还需要像 PPO 指南中那样训练奖励模型。我们将在 Anthropic-HH-RLHF 数据集上使用 REINFORCE 算法。

REINFORCE 训练#

在使用监督微调 (SFT) 微调 GPT 模型并按照前面部分中的说明训练奖励模型后，您可以开始使用 REINFORCE 对齐策略。

在 REINFORCE 训练期间，我们有三个模型相互交互，Aligner 在两个单独的作业中运行这些模型

  1. 策略网络：这是我们正在训练的模型，它应该从 SFT 模型开始。

  2. 奖励模型 (RM)：此模型接受提示和响应的组合作为输入，并生成单个标量值（称为奖励）。REINFORCE 算法旨在最大化此奖励。

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

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

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

启动服务器

#!/bin/bash
RM_NEMO_FILE="/path/to/trained_rm.nemo"
GPFS="/path/to/nemo-aligner-repo"

RESULTS_DIR="critic_results_dir"

cd ${GPFS}
export PYTHONPATH="${GPFS}:${PYTHONPATH}" \
&& export HYDRA_FULL_ERROR=1 \
&& python -u examples/nlp/gpt/serve_reward_model.py \
   trainer.num_nodes=1 \
   trainer.devices=8 \
   ++model.tensor_model_parallel_size=4 \
   rm_model_file=${RM_NEMO_FILE}

上面的示例在八个 GPU 和一个节点上启动奖励模型服务器。请务必根据您的模型大小和规模更改 trainer.devices、trainer.num_nodes。Aligner 将在任何规模上工作。此外，请务必调整 trainer.reinforce.inference_micro_batch_size 参数。此参数设置 REINFORCE actor 允许为每个 DP 排名发送给奖励的批次大小。

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

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

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

ACTOR_NEMO_FILE="/path/to/sft_checkpoint.nemo"
RESULTS_DIR="/path/to/actor_results_dir"

USE_FLASK=False
ACTOR_LR=1e-6
KL=0.01
NUM_ROLLOUTS=32
ACTOR_GBS=32
REWARD_PORT=5555
# Change this to the hostname of server hosting the reward model
host_reward="localhost"

cd ${GPFS}
export PYTHONPATH="${GPFS}:${PYTHONPATH}" \
&& export HYDRA_FULL_ERROR=1 \
&& mpirun -n 8 --allow-run-as-root python -u examples/nlp/gpt/train_gpt_reinforce_actor.py \
   "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.checkpoint_callback_params.save_top_k=1 \
   exp_manager.explicit_log_dir=\"${RESULTS_DIR}\" \
   trainer.reinforce.max_epochs=1 \
   trainer.reinforce.max_steps=313 \
   trainer.reinforce.val_check_interval=4 \
   trainer.num_nodes=1 \
   trainer.devices=8 \
   trainer.reinforce.trt_llm.enable=True \
   trainer.reinforce.trt_llm.reshard=True \
   trainer.reinforce.trt_llm.unload_engine_train=False \
   ++model.tensor_model_parallel_size=4 \
   ++model.reinforce.num_rollout_samples=${NUM_ROLLOUTS} \
   model.global_batch_size=${ACTOR_GBS} \
   model.micro_batch_size=1 \
   model.optim.lr=\"\\\$\{multiply:${ACTOR_LR},1.001\}\" \
   model.optim.sched.warmup_steps=0 \
   model.optim.sched.constant_steps=312 \
   model.optim.sched.min_lr=${ACTOR_LR} \
   model.optim.weight_decay=0.01 \
   model.reinforce.rollout_micro_batch_size=16 \
   model.reinforce.forward_micro_batch_size=16 \
   model.reinforce.val_rollout_micro_batch_size=8 \
   model.data.data_impl=jsonl \
   remote_rm.reward_model.ip=${host_reward} \
   remote_rm.reward_model.port=${REWARD_PORT} \
   ++model.reinforce.length_params.max_length=2048 \
   trainer.reinforce.initial_policy_kl_penalty="${KL}" \
   ++model.optim.bucket_cap_mb=200 \
   ++model.dist_ckpt_format=zarr \
   ++model.optim.overlap_grad_sync=False \
   ++model.optim.contiguous_grad_buffer=True \
   ++model.enable_nge=True \
   trainer.reinforce.batch_iterator.use_flask=${USE_FLASK} \
   trainer.reinforce.rollout_batch_seq_length=4096

上面的命令在一个节点上使用八个 GPU 启动初始服务器和 actor 服务器。

启动用于 REINFORCE 训练的两个服务器#

您可以使用 slurm 启动这两个作业，并通过以下方式使它们在完整的 REINFORCE 作业中协同工作

#!/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

NAME="reinforce"

# 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=reinforce_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"
REWARD_PORT=5567
CRITIC_CONFIG_PATH="${GPFS}/examples/nlp/gpt/conf"
CRITIC_CONFIG_NAME="inference_rm"

CONF_DIR="${GPFS}/examples/nlp/gpt/conf"
CONFIG_NAME="gpt_reinforce_actor"

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 examples/nlp/gpt/serve_reward_model.py \
   --config-path=${CRITIC_CONFIG_PATH} \
   --config-name=${CRITIC_CONFIG_NAME} \
   trainer.num_nodes=1 \
   trainer.devices=8 \
   ++model.tensor_model_parallel_size=4 \
   rm_model_file=${RM_NEMO_FILE} \
   inference.port=${REWARD_PORT}
EOF

srun --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"

NUM_ROLLOUTS=32
NORMALIZE="True"
ACTOR_LR="1e-6"
ACTOR_GBS=32
KL=0.01
USE_FLASK=False

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

ACTOR_NAME="${NAME}_actor"

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

read -r -d '' cmd_reinforce <<EOF
cd ${GPFS}
export PYTHONPATH="${GPFS}:${PYTHONPATH}" \
&& export HYDRA_FULL_ERROR=1 \
&& python -u examples/nlp/gpt/train_gpt_reinforce_actor.py \
   "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.checkpoint_callback_params.save_top_k=1 \
   exp_manager.explicit_log_dir=\"${RESULTS_DIR}\" \
   trainer.reinforce.max_epochs=1 \
   trainer.reinforce.max_steps=313 \
   trainer.reinforce.val_check_interval=4 \
   trainer.num_nodes=1 \
   trainer.devices=8 \
   trainer.reinforce.trt_llm.enable=True \
   trainer.reinforce.trt_llm.reshard=True \
   trainer.reinforce.trt_llm.unload_engine_train=False \
   ++model.tensor_model_parallel_size=4 \
   ++model.reinforce.num_rollout_samples=${NUM_ROLLOUTS} \
   model.global_batch_size=${ACTOR_GBS} \
   model.micro_batch_size=1 \
   model.optim.lr=\"\\\$\{multiply:${ACTOR_LR},1.001\}\" \
   model.optim.sched.warmup_steps=0 \
   model.optim.sched.constant_steps=312 \
   model.optim.sched.min_lr=${ACTOR_LR} \
   model.optim.weight_decay=0.01 \
   model.reinforce.rollout_micro_batch_size=16 \
   model.reinforce.forward_micro_batch_size=16 \
   model.reinforce.val_rollout_micro_batch_size=8 \
   model.data.data_impl=jsonl \
   remote_rm.reward_model.ip=${host_reward} \
   remote_rm.reward_model.port=${REWARD_PORT} \
   ++model.reinforce.length_params.max_length=2048 \
   trainer.reinforce.initial_policy_kl_penalty="${KL}" \
   ++model.optim.bucket_cap_mb=200 \
   ++model.dist_ckpt_format=zarr \
   ++model.optim.overlap_grad_sync=False \
   ++model.optim.contiguous_grad_buffer=True \
   ++model.enable_nge=True \
   trainer.reinforce.batch_iterator.use_flask=${USE_FLASK} \
   trainer.reinforce.rollout_batch_seq_length=4096
EOF

srun --mpi=pmix --het-group=1 -o $PPO_OUTFILE -e $PPO_ERRFILE --container-image=${CONTAINER} $MOUNTS bash -c "${cmd_reinforce}" &

wait

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

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

注意

请务必更改奖励模型参数 trainer.reinforce.inference_micro_batch_size，以使 trainer.reinforce.inference_micro_batch_size * DP size <= model.reinforce.rollout_micro_batch_size。

REINFORCE 结果#

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

在本页中