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

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:

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:

{
    "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.