huggingface / huggingface/trl

Support router-based multi-teacher (MOPD) distillation in DistillationTrainer

Open
#6,151 6 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.