Knowledge Distillation Customization
Learn how to train a smaller student model to mimic a larger teacher model using knowledge distillation (KD).
About
Knowledge Distillation transfers knowledge from a large teacher model to a smaller student model. During training, the student learns to match the teacher’s output probability distribution, producing a compact model that retains much of the teacher’s capability.
What you can achieve with KD:
- Compress models: Distill an 8B model into a 0.6B model for faster inference and lower deployment costs
- Reduce latency: Deploy a smaller model that responds faster while preserving quality
- Lower resource requirements: Serve a distilled model on fewer GPUs
KD vs SFT: Understanding the Trade-offs
Key Parameters
Workflow Overview
This tutorial follows a complete distillation pipeline:
- Fine-tune the teacher (SFT on the task dataset) so it learns the domain
- Establish a baseline by deploying the base student model and measuring ROUGE scores
- Distill into the student using the fine-tuned teacher’s soft targets
- Evaluate the distilled student and compare ROUGE scores against the baseline
When to choose KD:
- You have a high-quality large model and want a smaller, faster version
- Deployment latency or cost is a constraint
- The teacher and student share the same vocabulary (e.g., both are Qwen3 models)
When to choose SFT instead: Refer to the Full SFT tutorial when you want to train a model directly on labeled data without a teacher.
Prerequisites
Before starting this tutorial, ensure you have:
- Completed the Quickstart to install and deploy NeMo Helix locally
- Installed the Python SDK (PyPI wrapper:
pip install "nemo-helix[all]"; source checkout: runmake bootstrapfrom the repository root) - Installed evaluation dependencies:
- Two 80 GB GPUs with CUDA 13+ for the teacher’s full SFT. The distillation job and the deployments each use one GPU.
Quick Start
1. Initialize SDK
The SDK connects to the deployment in your current CLI context, using its base URL, default workspace, and credentials. To connect to another deployment, refer to Connect with CLI Contexts. These examples use the default workspace. If your current CLI context uses a different workspace, run nemo config set --workspace default first.
In environments without a CLI configuration file, such as Google Colab, set NHX_BASE_URL (and NHX_ACCESS_TOKEN for an authenticated deployment) before you create the client. NemoClient.from_config() reads both:
2. Prepare Dataset
Knowledge distillation uses the same dataset formats as SFT. We use the SQuAD dataset for both teacher training and distillation so that the teacher first learns the task, then transfers that knowledge to the student.
We convert each example to the chat messages format. The user turn is the same prompt the evaluation step sends, so the model is trained and queried through the same chat template, including the end-of-turn token that stops generation.
We also hold out a small test split for ROUGE evaluation at the end.
3. Secrets Setup
In this tutorial we use two Qwen3 models from Hugging Face:
- Teacher: Qwen/Qwen3-8B (8B parameters)
- Student: Qwen/Qwen3-0.6B (0.6B parameters)
Both models share the same tokenizer/vocabulary (required for knowledge distillation) and include a chat template for deployment with /chat/completions.
Hugging Face Authentication:
- Both models are public, so the Hugging Face token is optional
- For gated models, you must provide a Hugging Face token via the
token_secretparameter and accept the model’s terms on its Hugging Face page - Get your token from Hugging Face Settings (requires Read access)
4. Create Model FileSets and Model Entities
Knowledge distillation requires two model entities:
- Student model — the smaller model that will be trained (Qwen/Qwen3-0.6B)
- Teacher model — the larger model that provides soft targets (Qwen/Qwen3-8B)
Both are post-trained checkpoints, so the output model includes a chat template, which is required for the /chat/completions inference endpoint.
Phase 1: Fine-Tune the Teacher
5. Train Teacher with Full SFT
For best distillation results, fine-tune the teacher on the same dataset that will be used for distillation. This ensures the teacher has learned the task-specific knowledge that the student will inherit.
We train the 8B teacher with Full SFT on the SQuAD dataset, sharded across 2 GPUs. Activation checkpointing and a micro batch of 1 keep it within 80 GB per GPU. The global batch and learning rate follow the Automodel 7B full-SFT SQuAD recipe qwen2_5_7b_squad.yaml.
Phase 2: Establish Baseline (Base Student)
6. Deploy the Base Student Model
Before distillation, deploy the base student model (Qwen3-0.6B, without any fine-tuning) to establish a baseline ROUGE score. After distillation, we compare the distilled student against this baseline to measure improvement.
7. Generate Baseline Predictions on Test Set
8. Delete Baseline Deployment
Delete the baseline student deployment to free GPU resources for the distillation training job and subsequent distilled model deployment.
Phase 3: Distill into Student
9. Create Knowledge Distillation Job
Now create a distillation job that trains the 0.6B student using the fine-tuned 8B teacher’s output distribution. The KD loss settings and learning rate follow the Automodel reference recipe llama3_2_1b_kd.yaml. The model field specifies the student, and teacher_model references the trained teacher model entity from Phase 1.
GPU Requirements:
KD requires loading both student and teacher models, so plan GPU memory accordingly:
- 0.6B student + 8B teacher: 1 GPU (80 GB)
- 3B student + 8B teacher: 4 GPUs
- 8B student + 70B teacher: 8+ GPUs
Use teacher_precision="bf16" (default) to reduce teacher memory footprint.
10. Track Distillation Progress
The same progress display used for teacher training prints completion percentage and the latest loss, and charts loss by step once history is available.
Phase 4: Evaluate the Distilled Student Model
11. Deploy the Distilled Student Model
The output model has the same architecture as the 0.6B student—only its weights have been updated via distillation. It requires just 1 GPU to deploy.
12. Generate Student Predictions on Test Set
13. Compute ROUGE Scores
Compare the base student (before distillation) and the distilled student against the ground-truth reference completions using ROUGE metrics.
Interpreting ROUGE Scores:
What to expect:
- The base student (0.6B, no training) provides a lower bound since it has not seen the task data
- The distilled student (0.6B, KD) should significantly outperform the base student, demonstrating the knowledge transferred from the 8B teacher
- If the distilled student scores are not much higher than the baseline, try increasing
distillation_temperature, adjustingdistillation_ratio, or training for more epochs
Hyperparameters
For detailed information on all available hyperparameters, recommended values, and tuning guidance, refer to the Hyperparameter Reference.
Troubleshooting
Job fails during model download:
- Verify authentication secrets are configured (refer to Managing Secrets)
- For gated Hugging Face models, accept the license on the model page
- Check both
model(student) andteacher_modelURNs are correct - Ensure both model entities exist:
client.models.get_model(name=...)
Job fails with OOM (Out of Memory) error:
KD loads both models, so OOM is more likely than with SFT:
- First try: Use
teacher_precision="bf16"to reduce teacher memory - Still OOM: Reduce
micro_batch_sizeto 1 - Still OOM: Reduce
global_batch_sizeandmax_seq_length - Last resort: Increase
num_gpus_per_node
No chat template / /chat/completions fails:
- Use post-trained checkpoints (e.g.,
Qwen3-0.6B) instead of base checkpoints (Qwen3-0.6B-Base). Base checkpoints may not include a chat template in their tokenizer, so the output model would also lack one.
Distilled model quality is poor:
- Increase
distillation_temperature(try 2.0–5.0) to transfer more nuanced knowledge - Adjust
distillation_ratio—if dataset labels are high-quality, lower the ratio; if the teacher is strong, raise it - Increase
epochsormax_stepsfor more training - Verify teacher and student share the same vocabulary
Vocabulary mismatch error:
- Teacher and student must use the same tokenizer. Use models from the same family (e.g., Qwen3-0.6B + Qwen3-8B)
Deployment fails:
- Verify output model exists:
client.models.get_model(name=DISTILLED_STUDENT_NAME) - Check job logs:
client.jobs.list_job_logs(name=kd_job.job.name).items() - Check student deployment status:
client.models.get_deployment(name=student_deployment.name) - The distilled model has the same size as the student, so GPU requirements match the student model
Next Steps
- Monitor training metrics in detail
- Evaluate your fine-tuned model using the Evals service
- Learn about LoRA customization for resource-efficient fine-tuning
- Learn about Full SFT for direct supervised fine-tuning