重要提示

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

重要提示

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

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

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

使用知识蒸馏的监督式微调 (SFT)

知识蒸馏是一种技术，其中较小的（学生）模型从较大的（教师）模型学习。目标是从教师模型“提炼”信息到学生模型。与训练模型预测下一个 token 的标准 SFT 相比，知识蒸馏允许从教师模型到学生模型传递更精确的信息。与标准监督式微调相比，知识蒸馏有两个主要优点：（1）在更少的训练 token 中收敛，以及（2）提高的准确性。

知识蒸馏有很多变体。NeMo Aligner 支持训练学生模型以匹配教师模型的 top-K logits。在本教程中，我们将介绍使用微调后的 Nemotron 8B 聊天模型来微调 2B 学生模型。

注意

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

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

获取微调后的教师模型和预训练的学生模型

首先，我们必须先下载预训练的学生模型和微调后的教师模型。

2B GPT 学生模型

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

2. 使用以下命令将 NeMo 文件解压到文件夹中：mkdir student_checkpoint && tar -xvf GPT-2B-001_bf16_tp1.nemo -C student_checkpoint。

3. 然后运行脚本以将旧的 NeMo 检查点转换为 Megatron-Core 检查点。该脚本位于此处。

   python convert_gpt_nemo_to_mcore.py \
      --input_name_or_path ./student_checkpoint \
      --output_path ./2b_student.nemo
   

Nemotron 8B 教师模型

1. 将 Llama3-8B LLM 模型和分词器 下载到模型的文件夹中。您可以使用 Hugging Face CLI 来完成此操作

   huggingface-cli download nvidia/nemotron-3-8b-chat-4k-sft --local-dir teacher_checkpoint
   

完成这些步骤后，您应该拥有文件 2b_student.nemo 和 teacher_checkpoint/Nemotron-3-8B-Chat-4k-SFT.nemo，以便在 NeMo-Aligner 中使用。

注意

Megatron Core 模型使用 TransformerEngine 作为后端，它尝试查找高效的内核。但是，根据您的 GPU，它可能并不总是成功。如果您遇到与内核查找相关的错误，请在脚本顶部设置这些变量。

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

下载数据

在本示例中，我们将使用 OpenAssistant 数据集。使用以下脚本下载数据集并将其转换为聊天格式

python /opt/NeMo-Aligner/examples/nlp/data/steerlm/preprocess_openassistant_data.py --output_directory=data/oasst

缓存教师模型的 Logits

接下来，我们使用教师模型的 logits 扩充数据集。请注意，此代码将为每个示例生成降序排列的 top-K 教师 logits。为了本教程的目的，我们通过设置以下内容来保存教师模型的前四个 logits

top_k=4

在实践中，k 通常设置为更大的值，例如 100。

重要提示

未能以降序保存教师模型的 logits 可能会影响收敛。如果您选择使用与本示例中提供的脚本不同的脚本来计算教师模型的 logits，请确保生成的数据集仍然具有降序排列的教师模型 logits。

此步骤在 8 个 H100 80G GPU 上大约需要 20 分钟。

终端

要在终端上直接运行微调，请使用以下命令。为了成功执行，请确保 NeMo-Aligner 存储库设置为当前工作目录。

python examples/nlp/synthetic_data_gen/compute_topk_logits.py \
   trainer.num_nodes=1 \
   trainer.devices=8 \
   trainer.precision=bf16 \
   pretrained_checkpoint.restore_from_path=teacher_checkpoint/Nemotron-3-8B-Chat-4k-SFT.nemo \
   model.megatron_amp_O2=True \
   model.tensor_model_parallel_size=1 \
   data.chat=True \
   data.sample=True \
   data.num_workers=0 \
   data.data.max_seq_length=4096 \
   data.data.file_path=data/oasst/train.jsonl \
   data.data.add_eos=False \
   data.data.hf_dataset=True \
   top_k=4 \
   model.global_batch_size=16 \
   model.micro_batch_size=2 \
   start_from_idx=0 \
   end_at_idx=56439 \
   output_path=data/oasst/train_with_logits_0.jsonl

Slurm

要通过 Slurm 生成教师模型 logits，请运行以下命令

#!/bin/bash
#SBATCH -A <<YOUR ACCOUNT>>
#SBATCH -p <<<YOUR PARTITION>>>
#SBATCH -N 1
#SBATCH -t 2: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_ACTOR_NEMO_FILE="/path/to/your/mcore_gpt.nemo"

PROJECT=WANDB_PROJECT # if you want to use wandb

RESULTS_DIR="/path/to/result_dir"

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

CONTAINER=<<<CONTAINER>>> # use the latest NeMo Training container, Aligner will work there

MOUNTS="--container-mounts=MOUNTS" # mounts

read -r -d '' cmd <<EOF
echo "*******STARTING********" \
&& echo "---------------" \
&& echo "Starting..." \
&& cd ${GPFS} \
&& export PYTHONPATH="${GPFS}:${PYTHONPATH}" \
&& export HYDRA_FULL_ERROR=1 \
&& python examples/nlp/synthetic_data_gen/compute_topk_logits.py \
      trainer.num_nodes=\${SLURM_JOB_NUM_NODES} \
      trainer.devices=\${SLURM_NTASKS_PER_NODE} \
      trainer.precision=bf16 \
      pretrained_checkpoint.restore_from_path=teacher_checkpoint/Nemotron-3-8B-Chat-4k-SFT.nemo \
      model.megatron_amp_O2=True \
      model.tensor_model_parallel_size=1 \
      data.chat=True \
      data.sample=True \
      data.num_workers=0 \
      data.data.max_seq_length=4096 \
      data.data.file_path=data/oasst/train.jsonl \
      data.data.add_eos=False \
      data.data.hf_dataset=True \
      top_k=4 \
      model.global_batch_size=16 \
      model.micro_batch_size=2 \
      start_from_idx=0 \
      end_at_idx=56439 \
      output_path=data/oasst/train_with_logits_0.jsonl
EOF

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

您还可以通过替换以下行来为验证数据集生成教师模型 logits

data.data.file_path=data/oasst/train.jsonl \
end_at_idx=56399 \
output_path=data/oasst/train_with_logits_0.jsonl

为

data.data.file_path=data/oasst/val.jsonl \
end_at_idx=2937 \
output_path=data/oasst/val_with_logits_0.jsonl

注意

存储教师模型的 logits 可能会非常占用内存。为了避免在加载数据时内存溢出，数据被分块加载到内存中。上面的示例使用单个块。

要使用多个块，请多次运行代码，更改 start_from_idx 和 end_at_idx 索引以穷尽整个数据集

start_from_idx=${START_FROM_IDX} \
end_at_idx=${END_AT_IDX} \
output_path=data/oasst/train_with_logits_${CHUNK_INDEX}.jsonl

每次运行代码时，都会生成一个块。请注意，输出路径应附加块索引作为后缀。索引的范围预计为 0 到 num_chunks - 1。

微调学生模型

数据准备就绪后，您就可以微调学生模型了。此步骤在 8 个 H100 80G GPU 上大约需要 50 分钟。

终端

要在终端上直接运行微调，请使用以下命令。为了成功执行，请确保 NeMo-Aligner 存储库设置为当前工作目录。

python -u examples/nlp/gpt/train_gpt_knowledge_distillation.py \
   trainer.num_nodes=1 \
   trainer.devices=8 \
   trainer.precision=bf16 \
   trainer.knowledge_distillation.limit_val_batches=5 \
   trainer.knowledge_distillation.max_steps=100 \
   trainer.knowledge_distillation.val_check_interval=1000 \
   trainer.knowledge_distillation.save_interval=1000 \
   pretrained_checkpoint.restore_from_path=2b_student.nemo \
   model.megatron_amp_O2=True \
   ++model.tensor_model_parallel_size=1 \
   ++model.pipeline_model_parallel_size=1 \
   model.micro_batch_size=1 \
   model.global_batch_size=128 \
   model.optim.lr=1e-5 \
   model.optim.name=distributed_fused_adam \
   model.optim.weight_decay=0.01 \
   model.optim.sched.constant_steps=0 \
   model.knowledge_distillation.target_logits_scale=1.0 \
   model.knowledge_distillation.logits_scale=1.0 \
   model.knowledge_distillation.sft_loss_weight=0.4 \
   model.knowledge_distillation.kd_loss_weight=1 \
   model.knowledge_distillation.kd_loss=fwd_kl \
   "model.data.data_prefix={train: [data/oasst/train_with_logits_CHUNK_ID.jsonl], validation: [data/oasst/val_with_logits_CHUNK_ID.jsonl], test: [data/oasst/val_with_logits_CHUNK_ID.jsonl]}" \
   ++model.data.data_impl=chunked_jsonl \
   ++model.data.n_chunks=1 \
   ++"model.data.n_examples_per_chunk={train: 56440, validation: 2938, test: 2938}" \
   ++model.data.seq_length=4096 \
   model.data.splits_string=\'98,1,1\' \
   exp_manager.create_wandb_logger=True \
   exp_manager.wandb_logger_kwargs.name=sft_knowledge_distillation_70b_chat \
   exp_manager.wandb_logger_kwargs.project=sft_knowledge_distillation \
   exp_manager.explicit_log_dir=results/kd_log_dir \
   exp_manager.resume_if_exists=True \
   exp_manager.resume_ignore_no_checkpoint=True \
   exp_manager.checkpoint_callback_params.monitor=val_loss \
   exp_manager.checkpoint_callback_params.save_nemo_on_train_end=True

Slurm

要通过 Slurm 生成教师模型 logits，请运行以下命令

#!/bin/bash
#SBATCH -A <<YOUR ACCOUNT>>
#SBATCH -p <<<YOUR PARTITION>>>
#SBATCH -N 1
#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_ACTOR_NEMO_FILE="/path/to/your/mcore_gpt.nemo"

PROJECT=WANDB_PROJECT # if you want to use wandb

RESULTS_DIR="/path/to/result_dir"

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

CONTAINER=<<<CONTAINER>>> # use the latest NeMo Training container, Aligner will work there

MOUNTS="--container-mounts=MOUNTS" # mounts

read -r -d '' cmd <<EOF
echo "*******STARTING********" \
&& echo "---------------" \
&& echo "Starting training" \
&& cd ${GPFS} \
&& export PYTHONPATH="${GPFS}:${PYTHONPATH}" \
&& export HYDRA_FULL_ERROR=1 \
&& python -u /opt/NeMo-Aligner/examples/nlp/gpt/train_gpt_knowledge_distillation.py \
      trainer.num_nodes=\${SLURM_JOB_NUM_NODES} \
      trainer.devices=\${SLURM_NTASKS_PER_NODE} \
      trainer.precision=bf16 \
      trainer.knowledge_distillation.limit_val_batches=5 \
      trainer.knowledge_distillation.max_steps=100 \
      trainer.knowledge_distillation.val_check_interval=1000 \
      trainer.knowledge_distillation.save_interval=1000 \
      pretrained_checkpoint.restore_from_path=2b_student.nemo \
      model.megatron_amp_O2=True \
      ++model.tensor_model_parallel_size=1 \
      ++model.pipeline_model_parallel_size=1 \
      model.micro_batch_size=1 \
      model.global_batch_size=128 \
      model.optim.lr=1e-5 \
      model.optim.name=distributed_fused_adam \
      model.optim.weight_decay=0.01 \
      model.optim.sched.constant_steps=0 \
      model.knowledge_distillation.target_logits_scale=1.0 \
      model.knowledge_distillation.logits_scale=1.0 \
      model.knowledge_distillation.sft_loss_weight=0.4 \
      model.knowledge_distillation.kd_loss_weight=1 \
      model.knowledge_distillation.kd_loss=fwd_kl \
      "model.data.data_prefix={train: [data/oasst/train_with_logits_CHUNK_ID.jsonl], validation: [data/oasst/val_with_logits_CHUNK_ID.jsonl], test: [data/oasst/val_with_logits_CHUNK_ID.jsonl]}" \
      ++model.data.data_impl=chunked_jsonl \
      ++model.data.n_chunks=1 \
      ++"model.data.n_examples_per_chunk={train: 56440, validation: 2938, test: 2938}" \
      ++model.data.seq_length=4096 \
      model.data.splits_string=\'98,1,1\' \
      exp_manager.create_wandb_logger=True \
      exp_manager.wandb_logger_kwargs.name=sft_knowledge_distillation_70b_chat \
      exp_manager.wandb_logger_kwargs.project=sft_knowledge_distillation \
      exp_manager.explicit_log_dir=results/kd_log_dir \
      exp_manager.resume_if_exists=True \
      exp_manager.resume_ignore_no_checkpoint=True \
      exp_manager.checkpoint_callback_params.monitor=val_loss \
      exp_manager.checkpoint_callback_params.save_nemo_on_train_end=True
EOF

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

如果使用多个块运行，请相应地修改 data.n_chunks 和 data.n_examples_per_chunk。数据前缀（例如，data/oasst/train_with_logits_CHUNK_ID.jsonl）应保持不变。例如，如果 data.n_chunks=10 且 data.n_examples_per_chunk=100，我们应该有 10 个文件，名称为 data/oasst/train_with_logits_0.jsonl, ..., data/oasst/train_with_logits_9.jsonl。每个文件应包含 100 个示例。文件模板 data/oasst/train_with_logits_CHUNK_ID.jsonl 应传递给 modle.data.data_prefix，如上例所示。CHUNK_ID 将在数据加载时替换为 0 到 data.n_chunks-1 。

结果

下表展示了知识蒸馏的优势。使用 vanilla SFT 损失以及 SFT 损失和知识蒸馏损失的组合对预训练的 Nemotron-4 15B 模型进行了微调。知识蒸馏是使用 Nemotron-4 340B SFT 模型执行的。微调是使用数学/代码数据集执行的。

在吞吐量仅略有下降的情况下，知识蒸馏比 SFT 产生更高的准确率，同时需要更少的训练步骤才能收敛。

基础模型	训练目标	训练步数	MMLU (5-shot)	MMLU (0-shot)	HumanEval (0-shot)	MBPP (0-shot)	GSM8K (0-shot)	MATH (0-shot)
Nemotron 15B	SFT 损失	600,000	65.3	56.9	64.6	71.7	84.2	30.12
Nemotron 15B	KD + SFT 损失	420,000	65.3	57.3	70.1	73.3	85.2	35.84
Nemotron 15B	KD + SFT 损失	600,000	65.3	57.6	72	73.8	84.8	36.6