LLM Model Pruning and Knowledge Distillation with NVIDIA NeMo Framework

Model Pruning and Knowledge Distillation: A Powerful Combination for Smaller Language Models

Overview

Model pruning and knowledge distillation are powerful cost-effective strategies for obtaining smaller language models from an initial larger sibling. This tutorial focuses on creating a simple pipeline that prepares the dataset, fine-tunes the teacher on the WikiText-103-v1 dataset, and then prunes and distills the model to create a 4B model.

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:

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...")

Pruning and Distillation Pipeline

The pruning and distillation pipeline involves the following high-level steps (Figure 1):

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.
2. Distilling knowledge from teacher to student by using the 8B model as the teacher and the 4B pruned model as the student.

Results

Figures 6 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.

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.

FAQs

Q: What are the prerequisites for running this tutorial?
A: 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.

Q: How do I download the dataset?
A: Download the WikiText-103-v1 dataset and convert the train, test, and validation splits into JSONL files using the provided code or by running the introduction notebook.

Q: What are the steps involved in the pruning and distillation pipeline?
A: The pruning and distillation pipeline involves preparation, fine-tuning the teacher model, depth-pruning, width-pruning, and distilling knowledge from teacher to student.