LLM Model Pruning and Knowledge Distillation with NVIDIA NeMo Framework
Model pruning and knowledge distillation are powerful cost-effective strategies for obtaining smaller language models from an initial larger sibling.
Pruning:
Either drop layers (depth-pruning) or drop neurons, attention heads, and embedding channels (width-pruning).
Knowledge distillation:
Transfer knowledge from a large teacher model to a smaller student model, with the goal of creating a more efficient, smaller model that is faster and less resource-intensive to run.
The
How to Prune and Distill Llama-3.1 8B to an NVIDIA Llama-3.1-Minitron 4B Model
post discussed the best practices of using
large language models (LLMs)
that combine depth, width, attention, and MLP pruning with knowledge distillation–based retraining.
In this post, we provide a walk-through tutorial of the pruning and distillation pipeline in the NVIDIA NeMo framework on a simple dataset. This tutorial uses
Meta-Llama-3.1-8B
as the teacher model, with 4B being the target model size. We also visualize and discuss the training results.
Overview
This tutorial focuses on creating a simple pipeline that can prepare the dataset,  fine-tune the teacher on the
WikiText-103-v1
dataset, and then prune and distill the model to create the 4B model. The WikiText-103-v1 dataset is a collection of over 100M tokens extracted from a set of verified ‘Good’ and ‘Featured’ articles on Wikipedia. It is publicly available on Hugging Face.
In this tutorial, you are defining the pruning and distillation pipeline that involves the following high-level steps (Figure 1).
Figure 1. Steps from obtaining the dataset to creating the distilled 4B model
Preparation:
Download the dataset and convert to JSONL.
Preprocess by tokenizing the dataset.
Fine-tune the teacher model on the dataset.
Depth-prune the fine-tuned teacher model. The depth-pruned model is the starting point for the student network.
Width-prune the fine-tuned teacher model. The width-pruned model is the starting point for the student network.
Distilling knowledge from teacher to student by using the 8B model as the teacher and the 4B pruned model as the student.
To access the Jupyter notebooks for this tutorial, see the
/NVIDIA/NeMo
GitHub repo.
Prerequisites
You require access to at least eight NVIDIA GPUs with an individual memory of 80 GB, for example, eight H100-80GB or A100-80GB GPUs, and a Docker-enabled environment.
Follow the instructions in the project’s
README file
to install the NeMo framework, download the Meta-Llama-3.1-8B teacher model and get access to your Hugging Face access token.
Download the dataset
Download the WikiText-103-v1 dataset and convert the train, test, and validation splits into JSONL files using the following code or by running the
introduction notebook
:
# 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.")
Prepare the dataset
The pruning and the distillation scripts requires the data files to be preprocessed by tokenizing them using the meta-llama/Meta-Llama-3.1-8B tokenizer model to convert the data into a memory-map format. This can be done with the preprocessing script,
preprocess_data_for_megatron.py
in the NeMo framework.
Run the following script on the train split to prepare the dataset for pruning and distillation:
!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
Run the script on the test and the validation split as well. The
data preparation notebook
contains all the scripts to create the tokenized
wikitext_tokenized_{train/val/test}_text_document.{idx/bin}
files that can be used to fine-tune the teacher model.
Fine-tune the teacher model on the dataset
With the prepared dataset, perform a light fine-tuning procedure on the unpruned teacher model. This section shows the usage of the scripts rather than focusing on performance, so the fine-tuning setup is run with
GLOBAL_BATCH_SIZE
set to 128 and
STEPS
set to 30 to ensure a low training time.
Figure 2. Teacher fine-tuning
Run the
megatron_gpt_pretraining.py
script to correct the distribution shift across the original dataset on which the model was trained. Without correcting for the distribution shift, the teacher provides suboptimal guidance on the dataset when being distilled.
%%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} \
Running the script or executing the
teacher fine-tuning notebook
creates a fine-tuned teacher model.
Prune the fine-tuned teacher model to create a student
You can use two methods to prune the fine-tuned teacher model: depth-pruning and width-pruning.
Per the
technology report
, you can see that width-pruning generally outperforms depth-pruning in accuracy, but at the cost of increased inference latency. Choose to perform
depth-pruning
,
width-pruning
, or both methods based on these considerations.
Figure 3. Pruning the fine-tuned teacher model
Depth-prune the fine-tuned teacher model to create a student
In the first method, you depth-prune the model. To go from an 8B to a 4B model, prune the last 16 layers (layers 16-31). Run the
megatron_gpt_drop_layers.py
script to depth-prune the fine-tuned teacher model:
!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
Running this script or executing the
depth pruning notebook
results in the creation of a smaller checkpoint with the last 16 layers removed:
4b_depth_pruned_model.nemo
.
Width-prune the fine-tuned teacher model to create a student
In the second method, you width-prune the model. To go from an 8B to a 4B model, prune the model by reducing the MLP intermediate dimension and hidden size and by retraining the attention headcount and number of layers.
Run the
megatron_gpt_prune.py
script to width-prune the fine-tuned teacher model:
!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"
Running this script or executing the
width pruning notebook
results in the creation of a smaller width-pruned checkpoint:
4b_width_pruned_model.nemo
.
Distill knowledge from teacher into the student
The distillation process uses the fine-tuned model (8B) as the teacher model and the pruned model as the student model (4B) to distill to a smaller 4B model. Only logit loss is currently available in NeMo.
Figure 4. Distillation workflow
In this section, you distill the knowledge from the teacher model into both student models and compare them:
Distill the knowledge from the fine-tuned teacher to the depth-pruned student
Distill the knowledge from the fine-tuned teacher to the width-pruned student
Distill knowledge from the fine-tuned teacher into the depth-pruned student
Run the
megatron_gpt_distillation.py
script to distill knowledge from the teacher to the depth-pruned student model.
%%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}
Running this script or the
distilling depth-pruned student notebook
results in the creation of a distilled model:
depth_pruned_distilled_4b_model.nemo
.
Distill knowledge from the fine-tuned teacher into the width-pruned student
Run the
megatron_gpt_distillation.py
script to distill knowledge from the teacher to the width-pruned student model. Change the student model (
STUDENT
) and the directory where the distilled model will be saved (
FINAL_MODEL_PATH
) before running the script.
Running the
distilling width-pruned student notebook
results in the creation of a distilled model,
width_pruned_distilled_4b_model.nemo
.
Display the validation loss
Run the following code command or execute the
results notebook
to visualize the validation loss. Modify the path to the checkpoint before running the code example:
%load_ext tensorboard
%tensorboard --logdir "distill_trainings/megatron_llama_distill/" --port=6007
You can see the validation loss after running the distillation script over a
STEPS
value of 30, with the depth-pruned student and the width-pruned student in Figures 5 and 6, respectively.
Figure 5. Depth-pruned validation loss over 30 steps
Figure 6. Width-pruned validation loss over 30 steps
To configure this pipeline for your use case, run the scripts on a multi-node cluster with larger
GLOBAL_BATCH_SIZE, STEPS
and
VAL_INTERVAL
values to see an improvement in the validation loss.
Figures 7 and 8 show the validation loss decreasing when you run the training step in the distillation script over a
STEPS
value of 880 and a
GLOBAL_BATCH_SIZE
value of 2048 with the depth-pruned and width-pruned students, respectively.
Figure 7. Depth-pruned validation loss over 880 steps (with GBS=2048)
Figure 8. Width-pruned validation loss over 880 steps (with GBS=2048)
Conclusion
Pruning and distillation represent a significant advancement in the field of language model optimization. The ability to create smaller, more efficient models like the Llama-3.1-Minitron-4B in resource-constrained environments while preserving performance and without sacrificing substantial accuracy is a game changer for the AI industry.
The
Mistral-NeMo-Minitron-8B model
is developed using this approach and outperforms the Llama-3.1-8B model on a variety of
benchmarks
.
This approach reduces computational costs and energy consumption at inference time and also democratizes access to advanced NLP capabilities. This could revolutionize real-world applications in mobile devices, edge computing, and constrained resource settings. As these techniques continue to evolve, you can expect to see even more compact yet powerful language models, further expanding the reach of this technology across various industries.
For more information, see the following resources:
Jupyter notebooks
with pruning and distillation recipes
Compact Language Models via Pruning and Knowledge Distillation
research paper
LLM Pruning and Distillation in Practice: The Minitron Approach
with a discussion on performance metrics
How to Prune and Distill Llama-3.1 8B to an NVIDIA Llama-3.1-Minitron 4B Model
post introducing good practices surrounding pruning and distillation techniques
Mistral-NeMo-Minitron 8B Model Delivers Unparalleled Accuracy
post showing the performance benchmarks of the Mistral-NeMo-Minitron-8B model
使用 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)。
图 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，以确保缩短训练时间。
图 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。
从
技术报告
中可以看到，
宽度剪枝
的准确性通常优于
深度剪枝
，但代价是增加推理延迟。根据这些考虑因素，选择执行深度剪枝、宽度剪枝或这两种方法。
图 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 损失函数。
图 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 分别为深度剪枝学生和宽度剪枝学生。
图 5. Depth-pruned 验证损失超过 30 步
图 6、超过 30 步的宽度剪枝验证损失
要为您的用例配置此管道，请在具有更大
GLOBAL_BATCH_SIZE, STEPS
和
VAL_INTERVAL
值的多节点集群上运行脚本，以确保验证损失得到改善。
图 7 和图 8 显示，当您在蒸馏脚本中运行训练步骤时，在分别使用深度剪枝和宽度剪枝学生的情况下，
STEPS
值为 880 和
GLOBAL_BATCH_SIZE
值为 2048 时，验证损失会减少。
图 7、深度剪枝验证损失超过 880 步 (使用 GBS=2048 时)
图 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 模型的性能基准测试时，可提供无与伦比的准确性