End-to-end knowledge distillation

[tscholak]

🎯 Goal (What & Why)

Implement end-to-end knowledge distillation support, enabling training of a student model to mimic the outputs (logits) of a frozen teacher model.

🚀 Execution Plan

Step 1: What is the smallest working version?

• Use existing support for loading a frozen reference (teacher) model from Fast-LLM Issue #212.
• Add a knowledge distillation loss component (cross-entropy or KL-divergence) that compares logits from the teacher and student models.
• During training:
  • Compute logits from both the student and the frozen reference model.
  • Calculate and minimize the distillation loss.
  • Ensure the teacher model remains frozen and gradient-free.

Step 2: What additional optimizations are possible (but optional)?

• tbd

📌 Acceptance Criteria (Must-Haves for Completion)

• The distillation loss is integrated and functional.
• Teacher model logits can be seamlessly accessed during student training.
• Tested end-to-end distillation workflow is documented clearly.
• Benchmark demonstrating accuracy retention after distillation.

🛠️ Project Management

• Assign the project to the Fast-LLM project.
• Set the Estimate field (in days) in the GitHub project.
• Use the Size field to categorize the PR size (Small/Medium/Large).
• Assign an owner when opening the issue.

[nitsanluke]

Additional pointers from NeMo distillation implementation.

KL Loss
The above megatron_gpt_distillation.py makes use of modelopt.torch.distil model for the trainer loop (not essential for us at this point)