使用 NVIDIA NeMo 框架进行 LLM 模型剪枝和知识蒸馏

模型剪枝和知识蒸馏是功能强大且经济高效的策略，用于从最初较大的同级获得较小的语言模型。

剪枝：丢弃图层 (深度剪枝) 或丢弃神经元、注意力头和嵌入通道 (宽度剪枝)。
知识蒸馏： 将知识从大型教师模型转移到较小的学生模型，目标是创建更高效、更小、运行速度更快且资源密集型更低的模型。

在一篇“ 如何剪枝和蒸馏 Llama-3.1 8B ”博文中，讨论了使用大语言模型(LLM) 的最佳实践，该模型将深度、宽度、注意力和 MLP 剪枝与基于蒸馏的知识重新训练相结合。

在本文中，我们提供了一个关于 NVIDIA NeMo 框架中基于简单数据集的剪枝和蒸馏工作流的演练教程。本教程使用 Meta-Llama-3.1-8B 作为教师模型，目标模型大小为 4B。我们还会可视化并讨论训练结果。

概述

本教程重点介绍如何创建一个简单的工作流，用于准备数据集，针对 WikiText-103-v1 数据集对教师进行微调，然后对模型进行剪枝和蒸馏以创建 4B 模型。WikiText-103-v1 数据集包含从维基百科上一系列经过验证的“良好”和“精选”文章中提取的逾 100M 个令牌。它已在 Hugging Face 上公开发布。

在本教程中，您将定义涉及以下高级步骤的剪枝和蒸馏工作流 (图 1)。

A workflow diagram shows downloading the dataset, tokenizing, fine-tuning the 8B teacher dataset, pruning the teacher model, and distilling knowledge from teacher to student. — *图 1. 从获取数据集到创建蒸 4B 模型的步骤*

准备工作:
- 下载数据集并转换为 JSONL。
- 通过对数据集进行标记化预处理。
- 在数据集上微调教师模型。
- 深度剪枝微调的教师模型。深度剪枝模型是学生网络的起点。
- Width-prune 经过微调的教师模型。宽度剪枝模型是学生网络的起点。
通过将 8B 模型用作教师，将 4B 剪枝模型用作学生，将知识从教师提炼给学生。

要访问本教程中的 Jupyter 笔记本，请参阅 /NVIDIA/NeMo GitHub 存储库。

预备知识

您需要访问至少 8 个 NVIDIA GPUs（单个显存为 80 GB），例如 8 个 H100-80GB 或 A100-80GB GPUs，以及一个支持 Docker 的环境。

按照项目的 README 文件中的说明安装 NeMo 框架，下载 Meta-Llama-3.1-8B Instruct 模型，并获取 Hugging Face 访问令牌的访问权限。

下载数据集

下载 WikiText-103-v1 数据集，并使用以下代码或运行 introduction notebook ，将训练、测试和验证拆分转换为 JSONL 文件：

# Split into train, test and val files

import json
import os
from datasets import load_dataset

# Load the WikiText-103 dataset
dataset = load_dataset("wikitext", "wikitext-103-v1")

# Define the destination folder
data_folder = 'wikitext-data'
os.makedirs(data_folder, exist_ok=True)

# Define file paths and destination paths
file_paths = {
    'train': os.path.join(data_folder, 'wikitext-train.jsonl'),
    'validation': os.path.join(data_folder, 'wikitext-val.jsonl'),
    'test': os.path.join(data_folder, 'wikitext-test.jsonl')
}

# Function to save dataset split to a JSONL file
def save_to_jsonl(file_path, data):
    with open(file_path, 'w') as file:
        for item in data:
            file.write(json.dumps(item) + '\n')

# Define splits
splits = ["train", "validation", "test"]

# Save splits to JSONL files and calculate their sizes
for split in splits:
    if split in dataset:
        save_to_jsonl(file_paths[split], dataset[split])
    else:
        print(f"Split {split} not found in the dataset.")

准备数据集

剪枝和蒸馏脚本需要通过使用 meta-llama/Meta-Llama-3.1-8B 标记器模型对数据文件进行标记化来预处理数据文件，从而将数据转换为内存映射格式。这可以通过 NeMo 框架中的预处理脚本 preprocess_data_for_megatron.py 完成。

在 train split 中运行以下脚本，以准备用于剪枝和蒸馏的数据集：

!python /opt/NeMo/scripts/nlp_language_modeling/preprocess_data_for_megatron.py \
--input="./wikitext-data/wikitext-train.jsonl" \
--tokenizer-library='huggingface' \
--tokenizer-type='meta-llama/Meta-Llama-3.1-8B' \
--output-prefix=wikitext_tokenized_train \
--append-eod \
--workers=32

对测试和验证拆分运行脚本。数据准备 notebook 包含用于创建可用于微调 teacher model 的标记化 wikitext_tokenized_{train/val/test}_text_document.{idx/bin} 文件的所有脚本。

在数据集上微调教师模型

使用准备好的数据集，对未剪枝的教师模型执行微调过程。本节展示了脚本的用法，而非性能，因此运行微调设置时，将 GLOBAL_BATCH_SIZE 设置为 128，将 STEPS 设置为 30，以确保缩短训练时间。

A workflow diagram shows multiple steps: input token, embedding, transformer layers, LM head, Softmax, Logits, cross-entropy loss, and next token. Steps are marked as trainable or loss. — *图 2、教师微调*

运行 megatron_gpt_pretraining.py 脚本，以修正用于训练模型的原始数据集的分布偏移。在不修正分布偏移的情况下，教师会在提取数据集时提供次优指导。

%%bash 

export CUDA_DEVICE_MAX_CONNECTIONS=1

# Set path(s) if different:

MODEL="/workspace/llama-3_1-8b-nemo_v1.0/llama3_1_8b.nemo"

# Can change these to accommodate resources:

TENSOR_PARALLEL_SIZE=8
NODES=1
MICRO_BATCH_SIZE=4

# Don't change the following:

EXPERIMENT_DIR="distill_trainings"
EXPERIMENT_NAME="megatron_llama_ft"

DATA_TRAIN='wikitext_tokenized_train_text_document'
DATA_VAL='wikitext_tokenized_test_text_document'
DATA_TEST='wikitext_tokenized_val_text_document'

STEPS=30
GLOBAL_BATCH_SIZE=128

LOG_INTERVAL=1
VAL_INTERVAL=10
NUM_VAL_BATCHES=5

LR=1e-4
MIN_LR=1e-5
WARMUP_STEPS=2

cmd="torchrun --nproc-per-node=${TENSOR_PARALLEL_SIZE}"

${cmd} /opt/NeMo/examples/nlp/language_modeling/megatron_gpt_pretraining.py \
    --config-path /opt/NeMo/examples/nlp/language_modeling/conf/ \
    --config-name megatron_llama_distill.yaml \
    \
    name=${EXPERIMENT_NAME} \
    \
    exp_manager.exp_dir=${EXPERIMENT_DIR} \
    exp_manager.checkpoint_callback_params.save_top_k=1 \
    exp_manager.checkpoint_callback_params.save_nemo_on_train_end=True \
    \
    trainer.max_steps=${STEPS} \
    trainer.log_every_n_steps=${LOG_INTERVAL} \

运行脚本或执行教师微调 notebook 可创建经过微调的教师模型。

剪枝经过微调的教师模型以创建学生模型

您可以使用两种方法来剪枝经过微调的教师模型：depth-pruning 和 width-pruning。

从技术报告中可以看到，宽度剪枝的准确性通常优于深度剪枝，但代价是增加推理延迟。根据这些考虑因素，选择执行深度剪枝、宽度剪枝或这两种方法。

A diagram shows the iterative steps of training the LLM, estimating importance, ranking, trimming, and distilling. — *图 3、剪枝经过微调的教师模型*

对经过微调的教师模型进行深度剪枝，以创建一个学生模型

在第一种方法中，您可以对模型进行深度剪枝。要从 8B 到 4B 模型，请剪枝最后 16 层 (第 16 至 31 层)。运行 megatron_gpt_drop_layers.py 脚本以深度微调经过调优的教师模型：

!python -m torch.distributed.launch --nproc_per_node=8 \
     /opt/NeMo/examples/nlp/language_modeling/megatron_gpt_drop_layers.py \
     --path_to_nemo "./distill_trainings/megatron_llama_ft/checkpoints/megatron_llama_ft.nemo" \
     --path_to_save "/workspace/4b_depth_pruned_model.nemo" \
     --tensor_model_parallel_size 8 \
     --pipeline_model_parallel_size 1 \
     --gpus_per_node 8 \
     --drop_layers 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31

运行此脚本或执行深度剪枝 notebook 会导致创建较小的检查点，并删除最后 16 层：4b_depth_pruned_model.nemo。

Width-prune 经过微调的教师模型，以创建一个学生模型

在第二种方法中，您可以调整模型的宽度。要从 8B 模型升级到 4B 模型，请通过减少 MLP 中间维度和隐藏大小以及重新训练注意力头数和层数来剪枝模型。

运行 megatron_gpt_prune.py 脚本，以调整经过微调的教师模型的宽度：

!torchrun --nproc-per-node=8 /opt/NeMo/examples/nlp/language_modeling/megatron_gpt_prune.py \
     model.restore_from_path="./distill_trainings/megatron_llama_ft/checkpoints/megatron_llama_ft.nemo" \
     model.tensor_model_parallel_size=1 \
     model.pipeline_model_parallel_size=8 \
     +model.dist_ckpt_load_strictness=log_all \
     inference.batch_size=64 \
     trainer.num_nodes=1 \
     trainer.precision=bf16 \
     trainer.devices=8 \
     prune.ffn_hidden_size=9216 \
     prune.num_attention_heads=null \
     prune.num_query_groups=null \
     prune.hidden_size=3072 \
     export.save_path="/workspace/4b_width_pruned_model.nemo"

运行此脚本或执行宽度剪枝 notebook 会导致创建较小的宽度剪枝检查点：4b_width_pruned_model.nemo。

蒸馏知识从教师转化为学生模型

蒸馏过程将微调模型 (8B) 用作教师模型，将剪枝模型用作学生模型 (4B)，将蒸馏用作较小的 4B 模型。目前 NeMo 中只提供 logit 损失函数。

A workflow diagram shows classical knowledge distillation from teacher to student, with loss function from several layers of the transformer architecture. A student model with N layers is distilled from a teacher model with M layers. The student learns by minimizing a combination of embedding output loss, logit loss and transformer encoder specific losses mapped across student block S and teacher block T. — *图 4. 蒸馏工作流程*

在本节中，您将教师模型中的知识分为两个学生模型，并进行比较：

蒸馏从微调教师到深度剪枝学生的知识
蒸馏从微调教师到宽度剪枝学生的知识

蒸馏知识，从经过 fine-tuned 的教师到经过 depth-pruned 的学生模型

运行 megatron_gpt_distillation.py 脚本，将蒸馏知识从教师扩展到深度剪枝学生模型。

%%bash 

export CUDA_DEVICE_MAX_CONNECTIONS=1

# Can change these to accommodate resources:

TENSOR_PARALLEL_SIZE=8
NODES=1
MICRO_BATCH_SIZE=4

# Don't change the following:

EXPERIMENT_DIR="distill_trainings"
EXPERIMENT_NAME="megatron_llama_distill_depth_pruned_student"

TEACHER="${EXPERIMENT_DIR}/megatron_llama_ft/checkpoints/megatron_llama_ft.nemo"
STUDENT="/workspace/4b_depth_pruned_model.nemo"

FINAL_MODEL_PATH="${EXPERIMENT_DIR}/${EXPERIMENT_NAME}/checkpoints/depth_pruned_distilled_4b_model.nemo"

DATA_TRAIN='wikitext_tokenized_train_text_document'
DATA_VAL='wikitext_tokenized_test_text_document'
DATA_TEST='wikitext_tokenized_val_text_document'

STEPS=30
GLOBAL_BATCH_SIZE=128

LOG_INTERVAL=1
VAL_INTERVAL=10
NUM_VAL_BATCHES=5

LR=1e-4
MIN_LR=1e-5
WARMUP_STEPS=2

cmd="torchrun --nproc-per-node=${TENSOR_PARALLEL_SIZE}"

${cmd} /opt/NeMo/examples/nlp/language_modeling/megatron_gpt_distillation.py \
    name=${EXPERIMENT_NAME} \
    \
    exp_manager.exp_dir=${EXPERIMENT_DIR} \
    exp_manager.checkpoint_callback_params.save_top_k=1 \
    \
    trainer.max_steps=${STEPS} \
    trainer.log_every_n_steps=${LOG_INTERVAL} \
    trainer.val_check_interval=${VAL_INTERVAL} \
    trainer.limit_val_batches=${NUM_VAL_BATCHES} \
    +trainer.num_sanity_val_steps=0 \
    \
    trainer.precision=bf16 \
    trainer.devices=${TENSOR_PARALLEL_SIZE} \
    trainer.num_nodes=${NODES} \
    \
    "model.data.data_prefix={train:[1.0,$DATA_TRAIN],validation:[$DATA_VAL],test:[$DATA_TEST]}" \
    \
    model.restore_from_path=${STUDENT} \
    model.kd_teacher_restore_from_path=${TEACHER} \
    model.nemo_path=${FINAL_MODEL_PATH} \
    \
    model.tensor_model_parallel_size=${TENSOR_PARALLEL_SIZE} \
    model.sequence_parallel=True \
    model.micro_batch_size=${MICRO_BATCH_SIZE} \
    model.global_batch_size=${GLOBAL_BATCH_SIZE} \
    \
    model.optim.name=distributed_fused_adam \
    model.optim.lr=${LR} \
    model.optim.sched.min_lr=${MIN_LR} \
    model.optim.sched.warmup_steps=${WARMUP_STEPS}

运行此脚本或经过深度剪枝的提炼学生 notebook 会创建一个提炼模型：depth_pruned_distilled_4b_model.nemo。

蒸馏知识，从经过微调的教师到宽度剪枝的学生模型

运行 megatron_gpt_distillation.py 脚本，将蒸馏知识从教师传递到宽度剪枝的学生模型。在运行脚本之前，更改学生模型 (STUDENT) 和蒸馏模型的保存目录 (FINAL_MODEL_PATH)。

运行经宽度剪枝的提炼学生 notebook 会生成提炼模型 width_pruned_distilled_4b_model.nemo。

显示验证损失

运行以下代码命令或执行结果 notebook 以可视化验证损失。在运行代码示例之前，请修改检查点的路径：

%load_ext tensorboard
%tensorboard --logdir "distill_trainings/megatron_llama_distill/" --port=6007

当在 STEPS 值为 30 的情况下运行蒸馏脚本时，您可以看到验证损失，图 5 和图 6 分别为深度剪枝学生和宽度剪枝学生。

A plot shows the validation loss under 8 after running the training step in the distillation script for 30 steps with the depth-pruned student. — 图 5. Depth-pruned 验证损失超过 30 步

A plot shows the validation loss under 8 after running the training step in the distillation script for 30 steps with the width-pruned student. — *图 6、超过 30 步的宽度剪枝验证损失*

要为您的用例配置此管道，请在具有更大 GLOBAL_BATCH_SIZE, STEPS 和 VAL_INTERVAL 值的多节点集群上运行脚本，以确保验证损失得到改善。

图 7 和图 8 显示，当您在蒸馏脚本中运行训练步骤时，在分别使用深度剪枝和宽度剪枝学生的情况下，STEPS 值为 880 和 GLOBAL_BATCH_SIZE 值为 2048 时，验证损失会减少。

A plot shows the validation loss under 2.5 after running the training step in the distillation script with the depth-pruned model as the student. — *图 7、深度剪枝验证损失超过 880 步 (使用 GBS=2048 时)*

A plot shows the validation loss drop to under 2.5 after running the training step in the distillation script with the width-pruned model as the student. — *图 8、宽度剪枝验证损失超过 880 步 (使用 GBS=2048 时)*

结束语

剪枝和蒸馏代表了语言模型优化领域的重大进步。能够在资源受限的环境中创建更小、更高效的模型 (如 Llama-3.1-Minitron-4B)，同时保持性能且不牺牲大量准确性，这是 AI 行业的游戏规则变革。

Mistral-NeMo-Minitron-8B 模型是使用这种方法开发的，在各种基准测试中表现优于 Llama-3.1-8B 模型。

这种方法降低了推理时的计算成本和能耗，还普及了对高级 NLP 功能的使用。这可能会彻底改变移动设备、边缘计算和受限资源设置中的真实应用。随着这些技术的不断发展，您预计会看到更紧凑但强大的语言模型，进一步扩展这项技术的覆盖范围到各行各业。

有关更多信息，请参阅以下资源：

支持剪枝和蒸馏 recipes 的 Jupyter notebooks
通过剪枝和知识构建紧凑语言模型蒸馏 Compact Language Models via Pruning and Knowledge Distillation 研究论文
LLM 剪枝和蒸馏的实际应用：Minitron 方法与性能指标的讨论
如何剪枝和蒸馏 Llama-3.1 8B 到 NVIDIA Llama-3.1-Minitron 4B 模型的帖子，介绍了围绕剪枝和蒸馏技术的良好实践
Mistral-NeMo-Minitron 8B 模型在展示 Mistral-NeMo-Minitron 8B 模型的性能基准测试时，可提供无与伦比的准确性

使用 NVIDIA NeMo 框架进行 LLM 模型剪枝和知识蒸馏

概述

预备知识

下载数据集

准备数据集

在数据集上微调教师模型

剪枝经过微调的教师模型以创建学生模型

对经过微调的教师模型进行深度剪枝，以创建一个学生模型

Width-prune 经过微调的教师模型，以创建一个学生模型

蒸馏知识从教师转化为学生模型

蒸馏知识，从经过 fine-tuned 的教师到经过 depth-pruned 的学生模型

蒸馏知识，从经过微调的教师到宽度剪枝的学生模型

显示验证损失

结束语

相关资源

标签

关于作者

使用 NVIDIA NeMo 框架进行 LLM 模型剪枝和知识蒸馏

概述

预备知识

下载数据集

准备数据集

在数据集上微调教师模型

剪枝经过微调的教师模型以创建学生模型

对经过微调的教师模型进行深度剪枝，以创建一个学生模型

Width-prune 经过微调的教师模型，以创建一个学生模型

蒸馏知识从教师转化为学生模型

蒸馏知识，从经过 fine-tuned 的教师到经过 depth-pruned 的学生模型

蒸馏知识，从经过微调的教师到宽度剪枝的学生模型

显示验证损失

结束语

相关资源

标签

关于作者

相关文章

Mistral-NeMo-Minitron 8B 基础模型实现准确性巅峰

如何在 NVIDIA Llama-3.1-Minitron 4B 模型上修剪和提炼 Llama-3.1 8B

相关文章

借助 NVIDIA DriveOS LLM SDK 简化自动驾驶汽车应用的 LLM 部署

衡量 AI Guardrails 在生成式 AI 应用中的有效性和性能

聚焦：NAVER Place 利用 NVIDIA TensorRT-LLM 优化 SLM 基础的垂直服务

使用 DeepSeek-R1 NIM 构建具有专家推理功能的 AI 智能体

图像和视频理解的视觉语言模型提示工程实践指南