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).
- 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:
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:
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.
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.
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.
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:
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:
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.
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.
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.
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.
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
SAGECONNECT TECHNOLOGY GROUP CO., LIMITED
Sageconnect is a reliable provider of IT network system integration and solutions for global enterprise customers.
