Support router-based multi-teacher (MOPD) distillation in DistillationTrainer
- Dominant language
- Python
- Stars
- 19.3k
- Forks
- 3k
- Avg merge
- 1d 20h
- Merged PRs (30d)
- 194
Description
### Feature request
`DistillationTrainer` currently supports teacher-student distillation with a single teacher. A useful extension would be router-based multi-teacher distillation, where each example is routed to one teacher based on a dataset column such as `domain`, `task`, or `data_source`.
This would support workflows where a single student is trained from multiple specialist teachers, for example:
- math examples -> math teacher
- code examples -> code teacher
- general instruction examples -> general teacher
A possible API could be:
```python
DistillationConfig(
lmbda=1.0,
beta=1.0,
teacher_model_names_or_paths={
"math": "path-or-hub-id-to-math-teacher",
"code": "path-or-hub-id-to-code-teacher",
"general": "path-or-hub-id-to-general-teacher",
},
teacher_routing_column="domain",
)
```
Dataset example:
```python
{
"messages": [...],
"domain": "math",
}
```
Initial scope could be limited to:
> one example -> one route key -> one teacher
This would not require teacher aggregation, dynamic routing, teacher debate, or teacher updates. The main addition would be selecting the appropriate teacher per example before applying the existing distillation loss.
Relevant references:
- On-Policy Distillation: https://thinkingmachines.ai/blog/on-policy-distillation/
- MiMo-V2-Flash technical report, which introduces Multi-Teacher On-Policy Distillation: https://arxiv.org/abs/2601.02780
### Motivation
This would make `DistillationTrainer` more useful for mixed-domain post-training setups where different teacher models are strong in different domains.
Today, users can distill from a single teacher, but if they have separate specialist teachers for math, code, reasoning, safety, or general instruction, there does not seem to be a simple way to route each example to the appropriate teacher within the same training run.
Router-based multi-teacher distillation would allow users to train one deployable student model from multiple specialist teachers while keeping the first version simple and close to the existing `DistillationTrainer` flow.
### Your contribution
I would like to work on this and submit a PR if the maintainers are open to the feature.
I can start with a small scoped implementation that extends DistillationTrainer rather than adding a new trainer, preserves the existing single-teacher behavior, and adds routing from a dataset column to one configured teacher per example.
Contributor guide
Research direction
Start by locating DistillationTrainer and DistillationConfig, then trace the existing single-teacher selection and distillation-loss flow. Define the routing behavior around teacher_routing_column and teacher_model_names_or_paths, while preserving single-teacher behavior. Done means one example maps to one configured teacher and the relevant behavior is covered by the project's existing validation or training tests.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 38/100