重要提示

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

重要提示

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

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

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

通过 DPO、RPO 和 IPO 进行模型对齐#

NeMo Framework 通过 NeMo-Aligner 代码库支持高效的模型对齐。

NeMo-Aligner 中的所有算法都适用于来自 Megatron Core 的任何基于 GPT 的模型(在配置中,它具有 mcore_gpt=True)。为了本教程的目的,我们将使用新发布的具有 4096 序列长度的 2B GPT 模型,完成整个直接偏好优化 (DPO) 流程。相同的教程也适用于任何大小的 GPT 模型(例如 LLaMa3)。

使用 LoRA 的 DPO#

我们支持全参数 DPO 训练和 LoRA DPO 训练。在全参数 DPO 中,存在一个 actor 模型和一个参考模型。actor 模型使用参考模型初始化,并且是完全可训练的。参考模型被冻结,用于计算 KL 惩罚损失的对数概率(请参阅 DPO 论文)。对于基于 LoRA 的 DPO,actor 模型由参考模型加上 LoRA 权重初始化,其中只有 LoRA 权重是可训练的。因此,这使我们能够通过简单地启用或禁用 LoRA 来在 actor/参考模型之间切换。此外,无需存储两组 LLM 权重。

RPO 和 IPO 变体#

除了 vanilla DPO 算法外,我们还支持 DPO 算法的其他变体,包括 Identity Preference Optimization (IPO) 和 Reward-aware Preference Optimization (RPO)。

该算法通过 dpo.preference_loss 配置变量进行标识。我们支持三种基于距离度量的 RPO 算法:rpo_sq 用于平方距离,rpo_bwd_kl 用于伯努利后向 KL 散度,以及 rpo_fwd_kl 用于伯努利前向 KL 散度。

要使用 RPO 算法,每个数据集示例都应具有 chosen_rewardrejected_reward,它们可能来自人工标注员或奖励模型。如果数据中不存在 chosen_rewardrejected_reward,则使用 dpo.default_chosen_rewarddpo.default_rejected_reward

获取预训练模型#

首先,我们必须获取一个预训练模型以进行对齐。我们建议使用以下两个模型开始入门。本教程的其余部分将适用于任一模型,但出于演示目的,我们将使用较小的 2B 模型。

  1. 通过 wget http://hugging-face.cn/nvidia/GPT-2B-001/resolve/main/GPT-2B-001_bf16_tp1.nemo 获取 2B 检查点。

  2. 使用 mkdir model_checkpoint && tar -xvf GPT-2B-001_bf16_tp1.nemo -C model_checkpoint 将 NeMo 文件解压到文件夹。

  3. 运行脚本以将旧的 NeMo 检查点转换为 Megatron Core 检查点。该脚本位于此处
    python convert_nemo_gpt_to_mcore.py \
       --in-folder ./model_checkpoint \
       --out-file ./mcore_gpt.nemo
    
  1. Llama 3 8B LLM 模型和分词器下载到 models 文件夹中。您可以使用 Hugging Face CLI 执行此操作
    huggingface-cli download meta-llama/Meta-Llama-3-8B --local-dir /path/to/llama
    
  2. 将 LLaMa3 LLM 转换为 .nemo 格式。
    python /opt/NeMo/scripts/checkpoint_converters/convert_llama_hf_to_nemo.py \
        --input_name_or_path /path/to/llama --output_path /output_path/mcore_gpt.nemo
    

完成这些步骤后,您应该有一个文件 mcore_gpt.nemo,可在 NeMo-Aligner 中使用。

注意

Megatron Core 模型使用 TransformerEngine 作为后端,旨在查找高效的内核。但是,根据您拥有的 GPU,它可能并不总是能找到它们。如果您遇到与内核查找相关的错误,请考虑在脚本顶部设置以下变量。

export NVTE_MASKED_SOFTMAX_FUSION=0
export NVTE_FLASH_ATTN=0
export NVTE_FUSED_ATTN=0

此外,TransformerEngine 默认情况下是非确定性的。因此,使用相同参数进行的后续 DPO 运行将产生不同的结果,这对于参数扰动来说并不理想。有帮助的是,TransformerEngine 公开了一个标志,用于设置您是否要保证确定性的训练运行

export NVTE_ALLOW_NONDETERMINISTIC_ALGO=0

通过监督微调 (SFT) 教授指令跟随#

为了获得最佳的 DPO 训练性能,建议您从 SFT 模型而不是基础模型开始。有关如何在 Megatron GPT 模型上执行 SFT 的完整指南,请参阅 SFT 指南

DPO 模型训练#

准备数据集#

在运行核心 DPO 训练之前,您必须将训练和验证数据准备为 DPO 训练所需的格式。DPO 期望 .jsonl 文件,其中每行都是一个 JSON 字典,对应于单个完整的样本,如下所示

{"prompt": "Which year was the Magna Carta signed?", "chosen_response": "1215", "rejected_response": "I refuse to answer this question."}
{"prompt": "Please give me the name of a famous medieval painter.", "chosen_response": "Hieronymus Bosch", "rejected_response": "David Hockney"}

但是,请注意,大多数 Megatron GPT 模型都遵循必须遵循的严格格式模板。具体模板取决于 SFT 训练期间使用的模板。例如,许多 GPT 模型使用 extra_id 模板,应用该模板后,将要求您的数据格式如下所示

{"prompt": "<extra_id_0>System\n\n<extra_id_1>User\nWhich year was the Magna Carta signed?\n<extra_id_1>Assistant\n", "chosen_response": "1215\n<extra_id_1>", "rejected_response": "I refuse to answer this question.\n<extra_id_1>"}
{"prompt": "<extra_id_0>System\n\n<extra_id_1>User\nPlease give me the name of a famous medieval painter.\n<extra_id_1>Assistant\n", "chosen_response": "Hieronymus Bosch\n<extra_id_1>", "rejected_response": "David Hockney\n<extra_id_1>"}

始终遵循 SFT 训练期间用于 DPO 的 prompt-response 模板格式,因为不这样做将生成输出无用文本的模型。您应该创建一个 .jsonl 文件(采用上述格式)用于训练数据,并创建一个 .jsonl 文件用于验证数据。

您的 JSONL 文件必须包含至少与您计划在训练期间使用的全局批大小 (GBS) 一样多的样本。例如,如果 GBS = 64,请确保您的训练和验证文件都包含至少 64 个样本。使用样本数少于 GBS 的文件将导致崩溃。

使用 DPO 进行序列打包#

我们还支持使用 DPO 进行打包序列训练。序列打包是一种训练技术,其中多个训练示例被连接起来以创建一个更长的序列。这种方法消除了对填充的需求,并提高了 GPU 利用率。有关序列打包及其优点的详细概述,请参阅序列打包文档。本文档特别讨论了 SFT 的序列打包,但相同的优势也适用于 DPO。

DPO 数据集的打包作为 NeMo 和 NeMo-Aligner 中的预处理步骤完成。我们提供了一个 `script https://github.com/NVIDIA/NeMo-Aligner/blob/ashors/dpo-packing/examples/nlp/data/dpo/prepare_packed_dpo_dataset.py`_ 来打包您的 DPO 数据集。此脚本假定您已经准备好 DPO 格式的数据集。此脚本中运行三个主要步骤

  1. 运行 DPOModelDataset 中的在线处理代码。这包括诸如提示模板操作和分词之类的任务。结果是分词序列的数组,用索引表示。

  2. 连接选择的和拒绝的序列。

  3. 分词序列按长度分组,并运行打包算法。

您可以在此处阅读有关打包算法的更多信息。目前,支持 first_fit 的两种变体

  1. first_fit_decreasing:在应用 first-fit 算法之前,按降序对序列进行排序。它生成更优化的打包,但它倾向于将所有短序列保持在一起,这可能会对收敛产生影响。

  2. first_fit_shuffle:以随机顺序运行 first-fit。打包不太理想,但它保持数据集顺序随机。建议运行 first_fit_shuffle 并检查打包的序列长度。如果它们与目标长度相似(即高效打包),则使用 shuffle。否则,请尝试 first_fit_decreasing。

以下是运行打包脚本以准备 DPO 数据集的示例

python examples/nlp/data/dpo/prepare_packed_dpo_dataset.py \
   model.data.data_prefix=/path/to/training.jsonl \
   +model.encoder_seq_length=2048 \
   +tokenizer_path=/path/to/tokenizer/model \
   +output_dir=/path/to/output_folder \
   +pack_sizes=[4096] \
   +tokenizer_type=<huggingface or sentencpiece>
[  +packing_algorithm=first_fit_shuffle \  ]
[  ++model.seed=0                          ]

由于此脚本将选择的和拒绝的序列打包在一起,因此 pack_sizes 应始终至少是 model.encoder_seq_length 的两倍。当使用打包数据集运行训练时,model.encoder_seq_length 应设置为用于打包数据集的 packed_size

要在训练期间使用打包数据集,请将以下行添加到您的训练命令

++model.data.data_impl=packed_jsonl

使用序列打包运行训练时,需要记住以下几点

  1. 确保打包您的训练、验证和测试数据集。

  2. 序列打包只能以微批大小 1 运行。

  3. 序列打包通过 Transformer Engine 支持,因此请务必通过设置 ++model.transformer_engine=True 在配置中启用 transformer engine。

  4. 序列打包增加了每个全局批次处理的示例数量。尝试相应地缩放全局批大小,方法是将新的全局批大小设置为大约 unpacked_global_batch_size / avg_num_sequences_per_pack。平均每个包的序列数在 prepare_packed_dpo_dataset.py 完成后打印到 stdout。

开始训练#

一旦您的数据被处理成正确的格式,您就可以开始 DPO 训练了。您必须从预训练或 SFT 训练的模型开始。在本节中,我们将使用上一步中训练的 SFT 模型来训练 DPO 模型。对于以下部分的目的,我们假设您的训练 .jsonl 文件位于 /path/to/train_dpo_format.jsonl 中,而您的验证 .jsonl 文件位于 /path/to/valid_dpo_format.jsonl 中。

提示

如果您没有现成的 DPO 数据集,您可以生成一个玩具数据集以开始使用。以下是生成 NUM_EXAMPLES_TO_GENERATE 个示例的示例。确保此值大于 global_batch_size。

# Generates a dummy dataset in /path/to/train_dpo_format.jsonl /path/to/valid_dpo_format.jsonl

NUM_EXAMPLES_TO_GENERATE=200

mkdir -p /path/to
for i in $(seq 1 $NUM_EXAMPLES_TO_GENERATE); do
   cat <<EOF
{"prompt": "<extra_id_0>System\n\n<extra_id_1>User\n${i}*10=?\n<extra_id_1>Assistant\n", "chosen_response": "$((i * 10))\n<extra_id_1>", "rejected_response": "I refuse to answer this question.\n<extra_id_1>"}
EOF
done | tee /path/to/train_dpo_format.jsonl /path/to/valid_dpo_format.jsonl >/dev/null

对于以下参数,model.dpo.ref_policy_kl_penalty 对应于 DPO 论文中的 beta 参数。

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

export GPFS="/opt/NeMo-Aligner"
export TRAIN_DATA_PATH="/path/to/train_dpo_format.jsonl"
export VALID_DATA_PATH="/path/to/valid_dpo_format.jsonl"

python -u ${GPFS}/examples/nlp/gpt/train_gpt_dpo.py \
   trainer.num_nodes=1 \
   trainer.devices=8 \
   ++model.micro_batch_size=1 \
   ++model.global_batch_size=512 \
   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=dpo_training \
   exp_manager.wandb_logger_kwargs.name=dpo_training \
   exp_manager.explicit_log_dir=/results \
   ++trainer.dpo.max_epochs=1 \
   ++model.dpo.ref_policy_kl_penalty=0.1

要使用 Slurm 运行 DPO 模型训练,请使用以下脚本。该脚本使用 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 --gpus-per-node 8
#SBATCH --exclusive
#SBATCH --overcommit

export GPFS="/opt/NeMo-Aligner"
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=${GPFS}:${GPFS},${TRAIN_DATA_PATH}:${TRAIN_DATA_PATH},${VALID_DATA_PATH}:${VALID_DATA_PATH},${PRETRAINED_CHECKPOINT_NEMO_FILE}:${PRETRAINED_CHECKPOINT_NEMO_FILE}"

RESULTS_DIR="/path/to/result_dir"

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

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_gpt_dpo.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 \
   exp_manager.explicit_log_dir=${RESULTS_DIR} \
   exp_manager.create_wandb_logger=True \
   exp_manager.wandb_logger_kwargs.name=${NAME} \
   exp_manager.wandb_logger_kwargs.project=${PROJECT} \
   ++trainer.dpo.max_epochs=1 \
   ++model.dpo.ref_policy_kl_penalty=0.1
EOF

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

默认的 DPO 训练会调整所有参数。要使用 LoRA,我们可以设置 model.peft.peft_scheme=lora 并在 model.peft.lora_tuning 中使用不同的参数。请查看配置文件中的参数。

在 DPO 训练期间,将在 WandB 中记录多个指标,其中主要指标是 acc(表示模型的选定奖励超过拒绝奖励的百分比)。在这种情况下,reward 计算为模型对数概率与参考对数概率之间的差值,乘以 KL 惩罚(原始论文中的 beta),用于选定和拒绝的响应。在训练期间,acc 通常应增加,但如果其绝对值保持较低,请不要担心,因为它与最终的 MTBench 或 MMLU 分数无关。它应该只是总体上增加。

要监控的其他指标是 rewards_chosen_mean 和 rewards_rejected_mean,它们表示上述定义的 rewards 的平均值。虽然绝对值不一定至关重要,但重要的是 chosen_mean 始终如一地超过 rejected_mean。这些均值之间的差异越大越好。所有指标将按 WandB 中的 train/val/ 分组,分别表示该指标来自训练集还是验证集。

在 DPO 训练的理想超参数方面,很大程度上取决于您的 SFT 或基础/基础模型的特性。因此,没有适用于所有情况的通用参数。但是,以下列表简要概述了我们为各种模型大小扰动的超参数及其效果

  • global_batch_size:通常,我们发现,在所有其他参数相同的情况下,较低的 GBS 性能较差。对于我们训练的大多数模型,256 或 512 的 GBS 似乎是最佳点。

  • epochs:对训练数据大小高度敏感。我们建议您从 1 个 epoch 开始,然后再从此基础上增加。我们没有看到超过 3 个 epoch 的任何改进。

  • learning rate:我们测试了余弦退火,预热 10 步,然后缓慢衰减到恒定速率。该恒定速率应该相当低。我们看到 9e-7 和 5-e7 的最佳性能。

  • ref_policy_kl_penalty:我们通常看到较低的值(0.1、0.2、0.5 和 1.0)具有更好的性能。偶尔,高达 5.0 的值也有效。