模型剪枝和知識蒸餾是功能強大且經濟高效的策略，用于從最初較大的同級獲得較小的語言模型。
- 剪枝：丟棄圖層 (深度剪枝) 或丟棄神經元、注意力頭和嵌入通道 (寬度剪枝)。
- 知識蒸餾： 將知識從大型教師模型轉移到較小的學生模型，目標是創建更高效、更小、運行速度更快且資源密集型更低的模型。
在一篇“ 如何剪枝和蒸餾 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)。
- 準備工作:
- 下載數據集并轉換為 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 filesimport jsonimport osfrom datasets import load_dataset# Load the WikiText-103 datasetdataset = load_dataset("wikitext", "wikitext-103-v1")# Define the destination folderdata_folder = 'wikitext-data'os.makedirs(data_folder, exist_ok=True)# Define file paths and destination pathsfile_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 filedef save_to_jsonl(file_path, data): with open(file_path, 'w') as file: for item in data: file.write(json.dumps(item) + '\n')# Define splitssplits = ["train", "validation", "test"]# Save splits to JSONL files and calculate their sizesfor 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，以確保縮短訓練時間。
運行 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=8NODES=1MICRO_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=30GLOBAL_BATCH_SIZE=128LOG_INTERVAL=1VAL_INTERVAL=10NUM_VAL_BATCHES=5LR=1e-4MIN_LR=1e-5WARMUP_STEPS=2cmd="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。
從 技術報告 中可以看到， 寬度剪枝 的準確性通常優于 深度剪枝 ，但代價是增加推理延遲。根據這些考慮因素，選擇執行深度剪枝、寬度剪枝或這兩種方法。
對經過微調的教師模型進行深度剪枝，以創建一個學生模型
在第一種方法中，您可以對模型進行深度剪枝。要從 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 損失函數。
在本節中，您將教師模型中的知識分為兩個學生模型，并進行比較：
- 蒸餾從微調教師到深度剪枝學生的知識
- 蒸餾從微調教師到寬度剪枝學生的知識
蒸餾知識，從經過 fine-tuned 的教師到經過 depth-pruned 的學生模型
運行 megatron_gpt_distillation.py 腳本，將蒸餾知識從教師擴展到深度剪枝學生模型。
%%bash export CUDA_DEVICE_MAX_CONNECTIONS=1# Can change these to accommodate resources:TENSOR_PARALLEL_SIZE=8NODES=1MICRO_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=30GLOBAL_BATCH_SIZE=128LOG_INTERVAL=1VAL_INTERVAL=10NUM_VAL_BATCHES=5LR=1e-4MIN_LR=1e-5WARMUP_STEPS=2cmd="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 分別為深度剪枝學生和寬度剪枝學生。
要為您的用例配置此管道，請在具有更大 GLOBAL_BATCH_SIZE, STEPS 和 VAL_INTERVAL 值的多節點集群上運行腳本，以確保驗證損失得到改善。
圖 7 和圖 8 顯示，當您在蒸餾腳本中運行訓練步驟時，在分別使用深度剪枝和寬度剪枝學生的情況下，STEPS 值為 880 和 GLOBAL_BATCH_SIZE 值為 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 模型的性能基準測試時，可提供無與倫比的準確性
?
