重要提示

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

重要提示

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

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

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

通过拒绝采样进行模型对齐#

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

拒绝采样通常在监督式微调 (SFT) 之前进行。我们应首先遵循先决条件指南 和 SFT 指南。获得 SFT 模型后，我们还需要按照 PPO 指南 中的说明训练奖励模型。我们将在 Anthropic-HH-RLHF 数据集上使用拒绝采样算法。

拒绝采样训练#

在使用 SFT 微调 GPT 模型并按照上一节中的说明训练奖励模型后，您可以开始使用拒绝采样对齐策略。

在拒绝采样训练期间，我们有两个模型相互交互，NeMo-Aligner 在单独的作业中运行它们

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

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

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

启动奖励模型和批评者服务器#

要启动服务器

#!/bin/bash
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 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}

以上示例在 8 个 GPU 和 1 个节点上启动奖励模型服务器。请确保根据您的模型大小和规模更改 trainer.devices 和 trainer.num_nodes。NeMo-Aligner 将在任何规模上工作。另外，请确保调整 trainer.rs.inference_micro_batch_size 参数。此参数设置 RS Actor 允许每个 DP 排名发送给批评者的批次大小。

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

RS Actor 训练作业包含在需要时向所有服务器发出 HTTP 调用的主控制器。要启动 RS 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="/path/to/actor_results_dir"

ACTOR_LR=1e-6
NUM_ROLLOUTS=32
ACTOR_GBS=32
CRITIC_PORT=5555
host_critic="$(scontrol show hostnames=$SLURM_JOB_NODELIST_HET_GROUP_0 | head -n1)"

export PYTHONPATH="${GPFS}:${PYTHONPATH}" \
&& export HYDRA_FULL_ERROR=1 \
&& python -u examples/nlp/gpt/train_gpt_rs_actor.py \
   "model.data.data_prefix={train: [${TRAIN_DATA_PATH}], validation: [${VALID_DATA_PATH}], test: [${VALID_DATA_PATH}]}" \
   pretrained_checkpoint.restore_from_path=\"${PRETRAINED_ACTOR_NEMO_FILE}\" \
   exp_manager.checkpoint_callback_params.save_top_k=1 \
   exp_manager.explicit_log_dir=\"${RESULTS_DIR}\" \
   trainer.rs.max_epochs=1 \
   trainer.rs.max_steps=313 \
   trainer.rs.val_check_interval=4 \
   trainer.num_nodes=1 \
   trainer.devices=8 \
   ++model.tensor_model_parallel_size=4 \
   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.rs.num_rollout_samples=${NUM_ROLLOUTS} \
   model.rs.rollout_micro_batch_size=16 \
   model.rs.forward_micro_batch_size=16 \
   model.rs.val_rollout_micro_batch_size=8 \
   model.data.data_impl=jsonl \
   remote_rm.reward_model.ip=${host_critic} \
   remote_rm.reward_model.port=${CRITIC_PORT} \
   model.rs.num_rollouts_per_prompt=8 \
   model.rs.top_n_rollouts=1

以上命令在 1 个节点和 8 个 GPU 上启动 RS Actor 和初始策略服务器。

启动拒绝采样训练的两个服务器#

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

#!/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="2p_rs"

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

CONF_DIR="${GPFS}/examples/nlp/gpt/conf"
CONFIG_NAME="gpt_rs_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=${CRITIC_PORT}
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"

NUM_ROLLOUTS=16
NORMALIZE="True"
ACTOR_LR="1e-6"
ACTOR_GBS=16

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_rs <<EOF
cd ${GPFS} \
&& export PYTHONPATH="${GPFS}:${PYTHONPATH}" \
&& export HYDRA_FULL_ERROR=1 \
&& python -u examples/nlp/gpt/train_gpt_rs_actor.py \
   --config-path=${CONF_DIR} \
   --config-name=${CONFIG_NAME} \
   "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=\"${ACTOR_LOG_DIR}\" \
   trainer.rs.max_epochs=1 \
   trainer.rs.max_steps=313 \
   trainer.rs.val_check_interval=4 \
   trainer.num_nodes=1 \
   trainer.devices=8 \
   ++model.tensor_model_parallel_size=4 \
   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.rs.num_rollout_samples=${NUM_ROLLOUTS} \
   model.rs.rollout_micro_batch_size=2 \
   model.rs.forward_micro_batch_size=2 \
   model.rs.val_rollout_micro_batch_size=2 \
   model.data.data_impl=jsonl \
   remote_rm.reward_model.ip=${host_critic} \
   remote_rm.reward_model.port=${CRITIC_PORT} \
   model.rs.num_rollouts_per_prompt=8 \
   model.rs.top_n_rollouts=1
EOF

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

wait

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

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

注意

确保更改奖励模型参数 trainer.rs.inference_micro_batch_size，以使 trainer.rs.inference_micro_batch_size * DP size <= model.rs.rollout_micro_batch_size。

拒绝采样结果#

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

本页内容