LLM Pruning and Distillation in Practice: The Minitron Approach

• The method addresses the high cost of training model families from scratch when original pretraining data is unavailable for distillation.Existing pruning-and-distillation approaches assume access to the original pretraining dataset, an assumption that often fails for proprietary-data models.Source: Introduction · p. 1
• The authors introduce teacher correction to adapt the unpruned teacher to the distillation distribution and a downstream-task saliency criterion for depth pruning.They apply the strategy to Llama 3.1 8B and Mistral NeMo 12B, compressing them to 4B and 8B parameters, respectively.Source: Introduction · pp. 1–2
• MN-Minitron-8B outperforms similarly sized models across common language-modeling benchmarks, while both Llama 3.1-Minitron-4B variants show strong accuracy against larger and previous-generation models.The width-pruned Llama variant achieves better overall accuracy than the depth-pruned variant.Source: Introduction · p. 2
• The Llama 3.1-Minitron-4B depth- and width-pruned models achieve average TensorRT-LLM speedups of 2.7× and 1.8×, respectively, over the original Llama 3.1 8B.Source: Introduction · p. 2

Methodology

• The pipeline adjusts the teacher, prunes the model, and then uses distillation to recover accuracy in the compressed student.Source: Methodology · p. 2

Teacher Correction

• Teacher correction lightly fine-tunes the unpruned model on the distillation dataset before pruning and distillation, improving guidance when the original training data is unavailable.The authors also find that correcting the teacher in parallel with distillation performs on par with using a fully corrected teacher.Source: Teacher Correction · pp. 2–3, Figure 2 · p. 3

Pruning

• Structured pruning removes blocks or channels, ranking depth, neurons, attention heads, and embedding dimensions using activation sensitivity from forward passes on 1024 calibration samples.Importance is estimated simultaneously across the considered axes without backpropagation.Source: Pruning · p. 3
• Depth pruning ranks individual or contiguous layer removals using language-model validation loss or downstream-task accuracy, selecting Winogrande-based contiguous-layer pruning because validation loss can misrank downstream performance.The analysis also finds contiguous-layer removal preferable to dropping individual layers.Source: Pruning · p. 3
• After ranking, the method directly trims the corresponding MLP, attention, and normalization weight matrices and skips neural architecture search using prior architectural findings.Source: Pruning · p. 3

Retraining with Distillation

• Retraining compares conventional ground-truth training with knowledge distillation, using forward KL divergence between uncompressed-teacher and pruned-student logits.The teacher and student are the uncompressed and pruned models, respectively.Source: Retraining with Distillation · pp. 3–4

Pre-training

• The experiments use the Nemotron-4 curated continued-training dataset because the proprietary pretraining data for both source models is unavailable.The publicly available base models are used as starting points.Source: Pre-training · p. 4

Teacher Correction

• Teacher correction adapts Llama 3.1 8B and Mistral NeMo 12B to the distillation dataset using approximately 100B tokens and conservative optimization settings.The correction uses 120 warm-up steps, low learning rates, the original batch size, and the original decay schedule.Source: Teacher Correction · p. 4

Pruning

• The pruning recipe uses single-shot width pruning and contiguous depth pruning selected by Winogrande accuracy, with manually configured target architectures.The search phase is skipped, and the configurations are inspired by prior Minitron models.Source: Pruning · p. 4
• Llama 3.1 is compressed into width- and depth-pruned 4B variants, while Mistral NeMo 12B is reduced to an 8B width-pruned model.The Llama width variant reduces hidden and MLP dimensions, whereas the depth variant reduces depth from 32 to 16; the Mistral model reduces hidden and MLP dimensions.Source: Pruning · p. 5

Distillation

• Distillation minimizes forward KL divergence between teacher and student logits and omits the language-model cross-entropy loss.The unpruned model is the teacher and the pruned model is the student.Source: Pre-training · p. 4, Distillation · p. 5

Instruction Tuning

• Instruction tuning applies math and code SFT, instruction SFT, and two rounds of Reward-aware Preference Optimization to every model.The alignment procedure uses NeMo-Aligner.Source: Instruction Tuning · p. 5

Teacher Correction

• Teacher correction improves guidance relative to distilling directly from an uncorrected teacher, and correction can run in parallel with distillation to recover student accuracy.The parallel strategy performs on par with guidance from a fully corrected teacher.Source: Analysis · p. 5, Figure 4 · p. 6

Pruning and Distillation

• Pruning followed by distillation requires up to 40× fewer training tokens, using 380B instead of 15T tokens to produce MN-Minitron-8B.The comparison separates the benefits of pruning over random initialization and distillation over conventional language-model training.Source: Analysis · p. 5

Width vs. Depth Pruning

• Width pruning starts with lower loss and consistently outperforms depth pruning despite both Llama 3.1 variants having the same parameter count.The comparison is shown for the Llama 3.1-Minitron-4B variants.Source: Analysis · p. 5

Depth Pruning Metrics

• Non-contiguous layer removal can improve language-model validation loss, but downstream task performance instead favors removing layers 16 through 31 continuously.On Winogrande, per-layer importance selection gives 0.5 accuracy versus 0.595 for continuous removal, including after distillation retraining.Source: Analysis · pp. 5–6

Evaluation

• The evaluation spans language understanding, code generation, mathematics, commonsense reasoning, summarization, question answering, function calling, instruction following, and multiturn conversation.Base and aligned models are evaluated on different benchmark groups, including MMLU, HumanEval, MBPP, GSM8K, GPQA, BFCL, IFEval, and MT-Bench.Source: Evaluation · p. 6
• Teacher-correction convergence is compared between corrected and original teachers, including a setup where correction proceeds continuously alongside distillation.Figure 5 frames the parallel correction strategy as an alternative to first correcting the teacher and then distilling.Source: Figure 5 · p. 6

Base Models

• MN-Minitron-8B outperforms similarly sized models across common benchmarks, while Llama 3.1-Minitron-4B variants compare favorably with the 8B teacher and prior Minitron-4B.The width-pruned Llama variant performs better overall than the depth-pruned variant; MN-Minitron-8B uses 380B rather than 15T training tokens.Source: Base Models · p. 6

Instruct Models

• Aligned models outperform similarly sized variants on most evaluated benchmarks, with exceptions on HumanEval and MBPP and a MT-Bench deficit for Llama 3.1-Minitron-4B versus Gemma2.They are consistently better on MMLU, GSM8K, GPQA, IFEval, and BFCLv2.Source: Instruct Models · p. 6

Runtime Performance Analysis

• Training convergence comparisons include random initialization, random pruning, conventional language-model loss, and the proposed pruning-plus-distillation pipeline.Figure 6 presents these alternatives as the runtime-performance analysis setup.Source: Runtime Performance Analysis · pp. 6–7, Figure 6 · p. 7
• On a single H100 GPU, the depth-pruned 4B model achieves 2.7× average throughput over Llama 3.1 8B, the width-pruned model achieves 1.8×, and FP8 adds a 1.4× boost over BF16.The smaller model permits larger batch sizes in the tested runtime configurations.Source: Runtime Performance Analysis · p. 7

• The section signals that the evaluation yields interesting and surprising observations without specifying them in the supplied text.Source: Insights · p. 7

1. Teacher correction

• Teacher fine-tuning on the distillation dataset reduces LM validation loss by over 6% and does not affect pruning optimality.Teacher correction can also be performed in parallel with distillation.Source: §1 · p. 7
• Distillation reaches state-of-the-art post-pruning accuracy with 380B tokens instead of 15T, while width pruning benefits from retaining attention heads.The width-pruning strategy instead removes MLP intermediate dimensions and embedding channels.Source: §1 · p. 7

Mistral NeMo 12B to MN-Minitron-8B

• MN-Minitron-8B outperforms the teacher on GSM8K, rising from 55.7% to 58.5%, and HumanEval, rising from 23.8% to 36.2%.The reported improvement is likely influenced by the dataset, with retraining using distillation loss alone.Source: §1 · p. 7

Llama 3.1 8B to Llama 3.1-Minitron-4B

• Width pruning achieves higher Llama 3.1 compression accuracy than depth pruning, reaching 60.5% versus 58.7% on MMLU.Base-model GSM8K accuracy is 41.24% for width pruning versus 16.8% for depth pruning, although instruct tuning reduces the gap.Source: §1 · p. 7
• Depth pruning delivers a 2.7× throughput speedup over Llama 3.1 8B, compared with 1.7× for width pruning.For depth pruning, contiguous layer removal is more effective than non-contiguous importance-based pruning.Source: §1 · p. 7