Knowledge Distillation is a machine learning technique where a compact "student" model learns to replicate the behavior of a larger, more complex "teacher" model to achieve comparable performance with improved efficiency.

Model Optimizer's Distillation is a set of wrappers and utilities to easily perform Knowledge Distillation among teacher and student models. Given a pretrained teacher model, Distillation has the potential to train a smaller student model faster and/or with higher accuracy than the student model could achieve on its own.

This section focuses on demonstrating how to apply Model Optimizer to perform knowledge distillation with ease.

For Hugging Face models, please use the PyTorch docker image (e.g., nvcr.io/nvidia/pytorch:26.01-py3). Visit our installation docs for more information.

Also follow the installation steps below to upgrade to the latest version of Model Optimizer and install example-specific dependencies.

For Hugging Face models, install Model Optimizer with hf dependencies using pip from PyPI and install the requirements for the example:

pip install -U nvidia-modelopt[hf]
pip install -r requirements.txt

First obtain both a pretrained model to act as the teacher and a (usually smaller) model to serve as the student.

from transformers import AutoModelForCausalLM

# Define student & teacher
student_model = AutoModelForCausalLM.from_pretrained("student-model-id-or-path")
teacher_model = AutoModelForCausalLM.from_pretrained("teacher-model-id-or-path")

For HuggingFace models, ModelOpt provides KDTrainer, a drop-in replacement for HuggingFace's Trainer that handles the teacher forward pass and KD loss computation internally. Unlike the general-purpose Distillation API, KDTrainer does not call mtd.convert() and does not wrap the student in a DistillationModel — the student stays a plain HuggingFace model, and the teacher is kept on the trainer and forwarded explicitly during loss computation.

from modelopt.torch.distill.plugins.huggingface import KDTrainer

trainer = KDTrainer(
    student_model,
    training_args,
    distill_args={"teacher_model": teacher_model},  # criterion defaults to "logits_loss"
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
)

KDTrainer can be mixed in with other HuggingFace trainers (e.g. SFTTrainer) via normal Python multiple inheritance, as done in main.py:

class KDSFTTrainer(KDTrainer, SFTTrainer):
    pass

Since KDTrainer overrides compute_loss() to run the teacher forward pass and compute the KD loss, training is just the normal HuggingFace Trainer loop — no manual loss computation is required:

trainer.train()

Because the student is never wrapped in a DistillationModel, no mtd.export() step is needed — trainer.save_model() saves the student directly in its original HuggingFace format.

trainer.save_model(training_args.output_dir)

Loss criterion:

• mtd.LogitsDistillationLoss() - Standard KL-Divergence on output logits
• mtd.MGDLoss() - Masked Generative Distillation loss for 2D convolutional outputs
• mtd.MFTLoss() - KL-divergence loss with Minifinetuning threshold modification

Loss balancers:

• mtd.StaticLossBalancer() - Combines original student loss and KD loss into a single weighted sum (without changing over time)

Checkout the stand-alone distillation script in the examples/megatron_bridge/ for example scripts for KD with Megatron-Bridge which is generally more performant than the Hugging Face scripts.

Checkout the Knowledge Distillation example in the Megatron-LM repository.

In this e2e example we finetune Llama-3.2 models on the smol-smoltalk-Interaction-SFT dataset as a minimal example to demonstrate a simple way of integrating Model Optimizer's KD feature.

We replace normal supervised finetuning (SFT) of a Llama-3.2-1B base model by distilling information from Llama-3.2-3B-Instruct which has already been instruction-finetuned.

accelerate launch --config-file ./accelerate_config/fsdp2.yaml \
    main.py \
    --teacher_name_or_path 'meta-llama/Llama-3.2-3B-Instruct' \
    --student_name_or_path 'meta-llama/Llama-3.2-1B' \
    --output_dir ./llama3.2-distill \
    --max_length 2048 \
    --per_device_train_batch_size 4 \
    --per_device_eval_batch_size 8 \
    --max_steps 200 \
    --logging_steps 5

• 📅 Roadmap
• 📖 Documentation
• 🎯 Benchmarks
• 💡 Release Notes
• 🐛 File a bug
• ✨ File a Feature Request