Knowledge Distillation with NeMo AutoModel

This guide walks through fine-tuning a student LLM with the help of a larger teacher model using the kd (knowledge distillation) recipe.

In particular, we will show how to distill a 3B (meta-llama/Llama-3.2-3B) model into a 1B (meta-llama/Llama-3.2-1B) model.

What is Knowledge Distillation?

Knowledge distillation (KD) transfers the dark knowledge of a high-capacity teacher model to a smaller student by minimizing the divergence between their predicted distributions. The student learns from both the ground-truth labels (Cross-Entropy loss, CE) and the soft targets of the teacher (Kullback-Leibler loss, KD):

where $\alpha$ is the kd_ratio, $T$ is the softmax temperature, and $y$ contains the labels. The temperature-scaled teacher and student distributions are

Prepare the YAML Config

A ready-to-use example is provided at examples/llm_kd/llama3_2/llama3_2_1b_kd.yaml. Important sections:

• model – the student to be fine-tuned (1 B parameters in the example)
• teacher_model – a larger frozen model used for supervision (3 B)
• kd_ratio – blend between CE and KD loss
• temperature – softens probability distributions before KL-divergence
• peft – optional LoRA config (commented). Uncomment to train only a handful of parameters.

Feel free to tweak these values as required.

Example YAML

1
# Example config for knowledge distillation fine-tuning
2
# Run with:
3
#   automodel examples/llm_kd/llama3_2/llama3_2_1b_kd.yaml
4
5
recipe: KnowledgeDistillationRecipeForNextTokenPrediction
6
7
step_scheduler:
8
  global_batch_size: 32
9
  local_batch_size: 1
10
  ckpt_every_steps: 200
11
  val_every_steps: 100  # will run every x number of gradient steps
12
  num_epochs: 2
13
14
dist_env:
15
  backend: nccl
16
  timeout_minutes: 1
17
18
seed: 1111
19
20
# Student
21
model:
22
  _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained
23
  pretrained_model_name_or_path: meta-llama/Llama-3.2-1B
24
  torch_dtype: bf16
25
26
# Teacher
27
teacher_model:
28
  _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained
29
  pretrained_model_name_or_path: meta-llama/Llama-3.2-3B
30
  torch_dtype: bf16
31
32
checkpoint:
33
  enabled: true
34
  checkpoint_dir: checkpoints/
35
  model_save_format: safetensors
36
  save_consolidated: false
37
38
distributed:
39
  strategy: fsdp2
40
  dp_size: null
41
  tp_size: 1
42
  cp_size: 1
43
  pp_size: 1
44
  sequence_parallel: false
45
46
# PEFT can be enabled by uncommenting below – student weights will remain small
47
# peft:
48
#   _target_: nemo_automodel.components._peft.lora.PeftConfig
49
#   target_modules: '*_proj'
50
#   dim: 16
51
#   alpha: 32
52
#   use_triton: true
53
54
loss_fn:
55
  _target_: nemo_automodel.components.loss.masked_ce.MaskedCrossEntropy
56
57
# Knowledge-distillation hyper-params
58
kd_ratio: 0.5          # 0 → pure CE, 1 → pure KD
59
kd_loss_fn:
60
  _target_: nemo_automodel.components.loss.kd_loss.KDLoss
61
  ignore_index: -100
62
  temperature: 1.0
63
  fp32_upcast: true
64
65
# Optimizer
66
optimizer:
67
  _target_: torch.optim.Adam
68
  betas: [0.9, 0.999]
69
  eps: 1e-8
70
  lr: 1.0e-5
71
  weight_decay: 0
72
73
# Dataset / Dataloader
74
dataset:
75
  _target_: nemo_automodel.components.datasets.llm.squad.make_squad_dataset
76
  dataset_name: rajpurkar/squad
77
  split: train
78
79
dataloader:
80
  _target_: torchdata.stateful_dataloader.StatefulDataLoader
81
  collate_fn: nemo_automodel.components.datasets.utils.default_collater
82
  shuffle: false
83
84
validation_dataset:
85
  _target_: nemo_automodel.components.datasets.llm.hellaswag.HellaSwag
86
  path_or_dataset: rowan/hellaswag
87
  split: validation
88
  num_samples_limit: 64
89
90
validation_dataloader:
91
  _target_: torchdata.stateful_dataloader.StatefulDataLoader
92
  collate_fn: nemo_automodel.components.datasets.utils.default_collater

Current Constraints

• Pipeline parallelism requires pp_microbatch_size to equal pp_batch_size so each student microbatch uses the correct teacher logits.
• Student and teacher models must share the same tokenizer for now; support for different tokenizers will be added in the future.

For VLM knowledge distillation, use the KnowledgeDistillationRecipeForVLM recipe and the checked-in examples/vlm_kd/qwen3_5/qwen3_5_vl_4b_kd.yaml config.

Launch Training

Single-GPU Quick Run

$
# Runs on a single device of the current host
$
automodel examples/llm_kd/llama3_2/llama3_2_1b_kd.yaml

Multi-GPU (Single Node)

$
# Leverage all GPUs on the local machine
$
torchrun --nproc-per-node $(nvidia-smi -L | wc -l) \
>
    nemo_automodel/recipes/llm/kd.py \
>
    -c examples/llm_kd/llama3_2/llama3_2_1b_kd.yaml

Slurm Cluster

Copy the checked-in slurm.sub reference script, set its CONFIG variable to the knowledge-distillation YAML, adapt the cluster settings, and submit it directly:

$
cp slurm.sub kd.sub
$
# Edit kd.sub to set CONFIG=examples/llm_kd/llama3_2/llama3_2_1b_kd.yaml
$
# and configure the SBATCH directives, container, and mounts for your cluster.
$
sbatch kd.sub

See Run on a Cluster for the complete Slurm workflow.

Monitoring

Metrics such as loss, ce_loss, kd_loss, lr, and tps are logged to WandB when the corresponding section is enabled.

Checkpoints and Inference

• Checkpoints are written under the directory configured in the checkpoint.checkpoint_dir field at every ckpt_every_steps.
• The final student model is saved according to the checkpoint section (e.g., model_save_format: safetensors, consolidated weights if save_consolidated: final or save_consolidated: every).

The example config saves sharded Safetensors. Consolidate the latest checkpoint before loading it with the Hugging Face-compatible model API:

$
bash checkpoints/LATEST/model/consolidate.sh

Then tokenize the prompt, generate token IDs, and decode them:

1
import torch
2
from transformers import AutoTokenizer
3
4
import nemo_automodel as am
5
6
checkpoint = "checkpoints/LATEST/model/consolidated"
7
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
8
student = am.NeMoAutoModelForCausalLM.from_pretrained(
9
    checkpoint,
10
    force_hf=True,
11
    attn_implementation="sdpa",
12
)
13
student.eval()
14
15
inputs = tokenizer("Translate to French: I love coding!", return_tensors="pt").to(student.device)
16
with torch.inference_mode():
17
    output_ids = student.generate(**inputs, max_new_tokens=32)
18
19
print(tokenizer.decode(output_ids[0], skip_special_tokens=True))