重要提示
您正在查看 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_reward
和 rejected_reward
,它们可能来自人工标注员或奖励模型。如果数据中不存在 chosen_reward
和 rejected_reward
,则使用 dpo.default_chosen_reward
和 dpo.default_rejected_reward
。
获取预训练模型#
首先,我们必须获取一个预训练模型以进行对齐。我们建议使用以下两个模型开始入门。本教程的其余部分将适用于任一模型,但出于演示目的,我们将使用较小的 2B 模型。
通过
wget http://hugging-face.cn/nvidia/GPT-2B-001/resolve/main/GPT-2B-001_bf16_tp1.nemo
获取 2B 检查点。使用
mkdir model_checkpoint && tar -xvf GPT-2B-001_bf16_tp1.nemo -C model_checkpoint
将 NeMo 文件解压到文件夹。- 运行脚本以将旧的 NeMo 检查点转换为 Megatron Core 检查点。该脚本位于此处。
python convert_nemo_gpt_to_mcore.py \ --in-folder ./model_checkpoint \ --out-file ./mcore_gpt.nemo
- 将 Llama 3 8B LLM 模型和分词器下载到 models 文件夹中。您可以使用 Hugging Face CLI 执行此操作
huggingface-cli download meta-llama/Meta-Llama-3-8B --local-dir /path/to/llama
- 将 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
- 将 LLaMa3 LLM 转换为
完成这些步骤后,您应该有一个文件 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 格式的数据集。此脚本中运行三个主要步骤
运行
DPOModelDataset
中的在线处理代码。这包括诸如提示模板操作和分词之类的任务。结果是分词序列的数组,用索引表示。连接选择的和拒绝的序列。
分词序列按长度分组,并运行打包算法。
您可以在此处阅读有关打包算法的更多信息。目前,支持 first_fit
的两种变体
first_fit_decreasing
:在应用 first-fit 算法之前,按降序对序列进行排序。它生成更优化的打包,但它倾向于将所有短序列保持在一起,这可能会对收敛产生影响。
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 运行。
序列打包通过 Transformer Engine 支持,因此请务必通过设置 ++model.transformer_engine=True 在配置中启用 transformer engine。
序列打包增加了每个全局批次处理的示例数量。尝试相应地缩放全局批大小,方法是将新的全局批大小设置为大约
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 的值也有效。