重要提示

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

重要提示

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

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

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

通过自博弈微调 (SPIN) 进行模型对齐#

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

NeMo-Aligner 中的所有算法都适用于任何基于 GPT 的模型，这些模型来自 Megatron Core（在配置中，mcore_gpt=True）。在本教程中，我们将使用新发布的具有 4096 序列长度的 2B GPT 模型来完成整个 SPIN 流程。本教程也适用于任何大小的 GPT 模型（例如 LLaMa3）。

有关 SPIN 算法的详细信息，请参阅论文：https://arxiv.org/abs/2401.01335。

获取预训练模型#

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

  1. 通过 wget https://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 文件夹中。

  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 默认情况下是非确定性的，这意味着使用相同参数的 SPIN 后续运行将产生不同的结果，这对于参数扰动并不理想。有帮助的是，TransformerEngine 公开了一个标志，如果您想保证确定性的训练运行，可以设置该标志。

export NVTE_ALLOW_NONDETERMINISTIC_ALGO=0

SFT 与基础（Base）模型用于 SPIN 训练#

与 DPO 和 PPO 不同，SPIN 旨在基础（base）模型上运行——这些模型仅在自回归语言预测任务上训练，而不是在指令跟随任务上训练。但是，您也可以在已在基于指令的数据集上进行 SFT 的模型上运行 SPIN，类似于 DPO/PPO。这两种类型的模型都可以在 SPIN 中良好运行。如果您希望从监督微调模型而不是基础模型开始，请参阅我们关于如何在 Megatron GPT 模型上执行 SFT 的完整指南SFT 指南。

SPIN 模型训练#

SPIN 训练使用与 NeMo-Aligner SFT 训练器完全相同的数据集格式和文件。请参阅 SFT 的数据格式化部分，以了解 SPIN 所需的数据格式SFT 指南

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

对于以下参数，model.spin.ref_policy_kl_penalty 对应于 SPIN 论文中的 beta 参数，而 trainer.spin.max_iterations 对应于 T（每次迭代 trainer.spin.max_epochs 个 epoch）。

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

export GPFS="/path/to/nemo-aligner-repo"
export TRAIN_DATA_PATH="/path/to/train_spin_format.jsonl"
export VALID_DATA_PATH="/path/to/valid_spin_format.jsonl"

python -u ${GPFS}/examples/nlp/gpt/train_gpt_spin.py \
   trainer.num_nodes=1 \
   trainer.devices=8 \
   model.micro_batch_size=1 \
   model.global_batch_size=64 \
   pretrained_checkpoint.restore_from_path=/path/to/megatron_gpt_sft.nemo \
   "model.data.train_ds.file_path=${TRAIN_DATA_PATH}" \
   "model.data.validation_ds.file_path=${VALID_DATA_PATH}" \
   exp_manager.create_wandb_logger=false \
   exp_manager.wandb_logger_kwargs.project=spin_training \
   exp_manager.wandb_logger_kwargs.name=spin_training \
   exp_manager.explicit_log_dir=/results \
   trainer.spin.max_iterations=1 \
   trainer.spin.max_epochs=1 \
   model.spin.ref_policy_kl_penalty=0.1 \
   model.spin.length_params.max_length=2048 \
   model.data.train_ds.max_seq_length=4096

以下脚本使用 4 个节点，但您可以将节点计数更改为其他值。要使用 Slurm 运行 SPIN 模型训练

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

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

TRAIN_DATA_PATH="/path/to/train_spin_format.jsonl"
VALID_DATA_PATH="/path/to/valid_spin_format.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_spin.py \
   trainer.num_nodes=${SLURM_JOB_NUM_NODES} \
   trainer.devices=8 \
   pretrained_checkpoint.restore_from_path='${PRETRAINED_CHECKPOINT_NEMO_FILE}' \
   "model.data.train_ds.file_path=${TRAIN_DATA_PATH}" \
   "model.data.validation_ds.file_path=${VALID_DATA_PATH}" \
   model.micro_batch_size=1 \
   model.global_batch_size=64 \
   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.spin.max_iterations=1 \
   trainer.spin.max_epochs=1 \
   model.spin.ref_policy_kl_penalty=0.1 \
   model.spin.length_params.max_length=2048 \
   model.data.train_ds.max_seq_length=4096
EOF

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

在 SPIN 训练期间，几个指标将被记录到 WandB 以供您监控，主要是 acc（表示模型对真实响应的隐式奖励超过参考策略生成的响应的百分比）。在这种情况下，reward 的计算方法是模型 log probs 和参考 log probs 之间的差值，乘以真实响应和生成响应的 KL 惩罚（原始论文中的 beta）。在训练期间，acc 通常应呈上升趋势，但如果其绝对值保持较低，请不要担心，因为它与最终的 MTBench 或 MMLU 分数无关。它应该只是总体呈上升趋势。需要关注的其他指标是 rewards_actual_mean 和 rewards_generated_mean，它们表示上述 rewards 的平均值。同样，绝对值不一定那么重要，重要的是 actual_mean 应该随着时间的推移大于 generated_mean，并且差异越大越好。所有指标将按 WandB 中的 train/ 或 val/ 分组，分别表示该指标是来自训练集还是验证集。

注意

对于验证，我们仅计算 vanilla SFT 负对数似然损失，而不是使用正式的 SPIN 损失。因此，验证指标将仅包括 SFT NLL 损失。这种方法加快了验证过程，因为执行 SPIN 生成非常耗时，并且对于验证而言并非绝对必要。

当谈到 SPIN 训练的理想超参数时，很大程度上取决于您的 SFT（或基础/基础）模型和您的训练数据的特性，因此没有一个适用于所有情况的通用参数集。但是，以下简要概述了我们针对各种模型大小及其效果扰动的超参数

  • global_batch_size：SPIN 论文建议 7B 模型的 GBS 为 64，这与我们的发现一致——7B 模型的 GBS 越高，性能越差。对于更大的模型，您可以根据需要将 GBS 增加到 128 或 256，但建议从 64 作为基线开始。

  • iterations/epochs：SPIN 论文在 7B 模型上使用 iterations=3 和 epochs=2 进行训练，训练数据集大小为 20 万。使用与作者相同的基础模型，我们发现使用他们 20 万数据的 5 万子集时，iterations=1 和 epochs=1 效果更好。因此，我们建议从 iterations=1 开始，并根据需要在 MT-Bench/MMLU 上进行测试，增加到 2。

    此外，与 SPIN 论文不同，我们的实现目前不会将迭代 t-1 中生成的样本注入到迭代 t 中。这可以解释为什么我们没有看到迭代次数大于 1 时性能有任何提高。

  • learning rate：SPIN 论文建议从 5e-7 开始，并在最后一次迭代中退火至 1e-7。我们发现这通常效果很好；但是，我们也看到了恒定学习率为 4e-7 或 3e-7 的良好结果。

  • ref_policy_kl_penalty：这是一个正在进行研究的领域。SPIN 论文建议从 0.1 开始，并在最后一次迭代中增加到 5.0。我们发现 beta 为 0.1 在第一次迭代中效果良好，但随后的迭代往往会快速过拟合。提高 KL 惩罚似乎有所帮助，但不足以使 T > 1 检查点的性能优于 T <= 1。目前，我们建议将 KL 保留在 0.1，并且仅进行单次迭代训练。

本页内容