feat: Multi-Teacher On-Policy Distillation (MOPD)

docs/about/algorithms/index.md

@@ -10,6 +10,7 @@ NeMo RL supports multiple training algorithms for post-training large language m
 | [DAPO](dapo.md) | [DAPO Single Node](dapo.md#dapo-single-node) | [DAPO Multi-node](dapo.md#dapo-multi-node) |
 | [CISPO](cispo.md) | [CISPO Configuration](cispo.md#configuration) | [CISPO Async Lag-1 Recipe](cispo.md#async-lag-1-recipe) |
 | [On-policy Distillation](on-policy-distillation.md) | [Distillation Single Node](on-policy-distillation.md#on-policy-distillation-single-node) | [Distillation Multi-node](on-policy-distillation.md#on-policy-distillation-multi-node) |
+| [Multi-Teacher On-Policy Distillation (MOPD)](mopd.md) | — | [MOPD Configuration](mopd.md#configuration) |
 | [Supervised Fine-Tuning (SFT)](sft.md) | [SFT Single Node](sft.md#sft-single-node) | [SFT Multi-node](sft.md#sft-multi-node) |
 | [DPO](dpo.md) | [DPO Single Node](dpo.md#dpo-single-node) | [DPO Multi-node](dpo.md#dpo-multi-node) |
 | [PPO](ppo.md) | [PPO Single Node](ppo.md#ppo-single-node) | [PPO Multi-node](ppo.md#ppo-multi-node) |
@@ -25,6 +26,7 @@ dapo
 cispo
 ppo
 on-policy-distillation
+mopd
 sft
 dpo
 rm

docs/about/algorithms/mopd.md

@@ -0,0 +1,148 @@
+# Multi-Teacher On-Policy Distillation (MOPD)
+
+Multi-Teacher On-Policy Distillation (MOPD) distills one or more teacher models
+into the policy by replacing GRPO's reward-based advantage with a token-level
+distillation advantage ([MiMo-V2-Flash Technical Report](https://arxiv.org/abs/2601.02780)).
+MOPD runs on async GRPO and collects rollouts through NeMo Gym, so the agent
+loop drives multi-turn / multi-step interaction. Each token of the resulting
+student rollout is scored by a teacher, and the policy is updated to close the
+gap with the teacher.
+
+Unlike the teacher-logit knowledge distillation in
+[On-policy Distillation](on-policy-distillation.md) (`run_distillation.py`), MOPD
+runs on top of the GRPO trainer: it is selected with `adv_estimator: opd` and
+serves teachers from dedicated, non-colocated worker groups during async
+collection.
+
+## Advantage
+
+For each token `t`, the distillation advantage is the stop-gradient
+teacher-minus-student log-probability gap:
+
+```
+Â_t = sg[ log π_teacher(t) − log π_student(t) ]
+```
+
+`log π_student` is the policy's `prev_logprobs` and `log π_teacher` is computed
+by the teacher worker group at collection time. Maximizing this advantage is
+reverse-KL minimization — it pushes the student toward the teacher's token
+distribution — but, unlike forward-KL logit distillation, it needs only the
+teacher's log-probability for the *sampled* token rather than the full
+vocabulary distribution.
+
+The advantage is applied only to trained (assistant) tokens via the loss mask;
+tool / environment tokens contribute zero. Because the advantage subtracts a
+real `prev_logprobs`, MOPD requires the student log-probabilities to actually be
+computed — see [Configuration](#configuration).
+
+## Configuration
+
+Enable MOPD in two places: select the advantage estimator and add the
+`on_policy_distillation` block.
+
+```yaml
+grpo:
+  # MOPD runs on async GRPO with NeMo Gym rollouts.
+  async_grpo:
+    enabled: true
+  adv_estimator:
+    name: opd
+  # OPD subtracts a real prev_logprobs, so it must not be skipped.
+  seq_logprob_error_threshold: 2.0
+
+loss_fn:
+  # REINFORCE form (drop the PPO probability-ratio clipping); on-policy
+  # correction is handled by the ICE-POP gate below instead.
+  disable_ppo_ratio: true
+  # ICE-POP hard gate: zero tokens whose train/inference importance-sampling
+  # weight falls outside bounds, correcting async off-policy drift.
+  use_importance_sampling_correction: true
+  truncated_importance_sampling_type: icepop
+  # Teacher distillation is the entire learning signal — no reference-policy KL.
+  reference_policy_kl_penalty: 0.0
+
+on_policy_distillation:
+  enabled: true
+  # Map each NeMo Gym agent name to a teacher checkpoint.
+  teacher_model_by_agent_name:
+    default_teacher: Qwen/Qwen3-1.7B
+  # Agents not present in the map fall back to this alias (must be a mapped key).
+  default_teacher_alias: default_teacher
+  # If true, an unmapped agent raises instead of falling back.
+  strict_agent_name_match: false
+  # Aliases that share one checkpoint reuse a single teacher worker group.
+  deduplicate_shared_teacher_checkpoints: true
+  non_colocated_teachers:
+    enabled: true
+    # Resourcing for each teacher worker group.
+    default_teacher_cfg:
+      tensor_model_parallel_size: 2
+      pipeline_model_parallel_size: 1
+      context_parallel_size: 1
+      num_nodes: 1
+      gpus_per_node: 8
+      precision: bf16
+      micro_batch_size: 1
+    # Optional per-alias overrides on top of default_teacher_cfg.
+    teacher_overrides: {}
+```
+
+> [!NOTE]
+> Teachers run the Megatron backend in inference-only mode. A DTensor-configured
+> policy is rejected for the teacher; PEFT / draft modules are stripped so
+> adapters are never attached to the frozen teacher; and teachers run
+> unquantized (a policy `quant_cfg` is ignored, with a warning).
+
+> [!NOTE]
+> `adv_estimator: opd` fails fast at setup if the config would zero
+> `prev_logprobs` (`loss_fn.force_on_policy_ratio: true` with no
+> `grpo.seq_logprob_error_threshold`), because the advantage would silently
+> degrade to `teacher_logprobs − 0`.
+
+### Teacher routing
+
+Each rollout sample carries its NeMo Gym `agent_ref`. At collection time the
+agent name is resolved to a teacher alias (`teacher_model_by_agent_name`, falling
+back to `default_teacher_alias`), samples are grouped by teacher, and each group
+is scored by exactly one teacher — there is no ensemble averaging across
+teachers. When several aliases map to the same checkpoint,
+`deduplicate_shared_teacher_checkpoints` collapses them onto a single worker
+group so they share GPUs.
+
+### Resourcing
+
+Non-colocated teachers each get their own Ray cluster on dedicated GPUs (they
+are queried every rollout group, so time-sharing with the policy/generation
+would serialize and destroy the async overlap). Their nodes are reserved from
+the policy's budget: with `total_nodes` total, the teacher groups take
+`sum(num_nodes)` and the policy uses the remainder (setup fails if nothing is
+left for the policy). Deduplicated teachers share one group's nodes.
+
+For example, the reference 3-node recipe lays out: 1 node policy (student,
+trainable) + 1 node vLLM generation (frozen) + 1 node teacher (frozen). Ten
+distinct teachers at 1 node each would instead add 10 nodes on top of the
+policy and generation nodes.
+
+## Running MOPD
+
+MOPD collects rollouts through NeMo Gym, so use the NeMo Gym GRPO entrypoint
+with an MOPD recipe. The checked-in recipe uses placeholder dataset paths;
+override them for your local data:
+
+```sh
+uv run examples/nemo_gym/run_grpo_nemo_gym.py \
+  --config examples/configs/recipes/llm/mopd-qwen3-1.7b-3n8g-megatron-pack.yaml \
+  data.train.data_path=/path/to/train.jsonl \
+  data.val.data_path=/path/to/val.jsonl
+```
+
+The reference recipe self-distills `Qwen/Qwen3-1.7B` (student == teacher) across
+3 nodes (1 policy + 1 vLLM + 1 teacher) with sequence packing enabled. Because
+student and teacher are identical, the OPD loss stays near zero — it is a
+correctness smoke test, not a demonstration of distillation gains.
+
+## References
+
+- LLM-Core Xiaomi, *MiMo-V2-Flash Technical Report*, which introduces the
+  multi-teacher on-policy distillation paradigm:
+  [arxiv.org/abs/2601.02780](https://arxiv.org/abs/2601.02780)

docs/index.md

@@ -62,7 +62,7 @@ Learn about DTensor and Megatron Core training backends, their capabilities, and
 :link: about/algorithms/index
 :link-type: doc
 
-Discover supported algorithms including GRPO, PPO, SFT, DPO, RM, and on-policy distillation with detailed guides and examples.
+Discover supported algorithms including GRPO, PPO, SFT, DPO, RM, on-policy distillation, and multi-teacher on-policy distillation (MOPD) with detailed guides and examples.
 :::
 
 :::{grid-item-card} {octicon}`graph` Evaluation

examples/configs/grpo_math_1B.yaml

@@ -458,3 +458,11 @@ data_plane:
   local_buffer_size:   68719476736   # 64 GiB  — used when backend == "mooncake_cpu"
   # observability:                    # NotRequired
   #   enabled: false
+
+# Multi-Teacher On-Policy Distillation (MOPD): distills from one or more teacher
+# models into the policy via token-level teacher-minus-student logprob advantages,
+# served by non-colocated teacher worker groups (OPD advantage estimator +
+# nemo_gym). null = disabled (default). See
+# examples/configs/recipes/llm/mopd-qwen3-1.7b-3n8g-megatron-pack.yaml for a full
+# enabled example.
+on_policy_distillation: null

examples/configs/recipes/llm/mopd-qwen3-1.7b-3n8g-megatron-pack.yaml

@@ -0,0 +1,144 @@
+defaults: ../../grpo_math_1B.yaml
+grpo:
+  num_prompts_per_step: 8
+  num_generations_per_prompt: 4
+  num_val_generations_per_prompt: 1
+  max_num_steps: 5
+  val_period: 1000
+  overlong_filtering: true
+  max_val_samples: null
+  val_batch_size: 32
+  async_grpo:
+    enabled: true
+  adv_estimator:
+    name: opd
+  seq_logprob_error_threshold: 2.0
+loss_fn:
+  reference_policy_kl_penalty: 0.0
+  kl_input_clamp_value: null
+  kl_output_clamp_value: null
+  ratio_clip_max: 0.28
+  use_on_policy_kl_approximation: true
+  disable_ppo_ratio: true
+  use_importance_sampling_correction: true
+  truncated_importance_sampling_ratio: 5.0
+  truncated_importance_sampling_ratio_min: 0.2
+  truncated_importance_sampling_type: icepop
+checkpointing:
+  enabled: false
+  checkpoint_dir: results/mopd-qwen3-1.7b-3n8g-megatron-pack
+  metric_name: val:total_reward/mean
+  keep_top_k: 1
+  save_period: 1000
+  checkpoint_must_save_by: 00:03:40:00
+  save_optimizer: false
+policy:
+  model_name: Qwen/Qwen3-1.7B
+  train_global_batch_size: 32
+  train_micro_batch_size: 1
+  generation_batch_size: 64
+  logprob_batch_size: 1
+  max_total_sequence_length: 32768
+  logprob_chunk_size: 2048
+  dtensor_cfg:
+    enabled: false
+  megatron_cfg:
+    enabled: true
+    activation_checkpointing: true
+    bias_activation_fusion: false
+    tensor_model_parallel_size: 2
+    sequence_parallel: true
+    defer_fp32_logits: true
+    optimizer:
+      lr: 3.0e-06
+      min_lr: 3.0e-06
+      weight_decay: 0.0
+    scheduler:
+      lr_decay_iters: null
+      lr_warmup_iters: 10
+      lr_warmup_init: 3.0e-07
+    distributed_data_parallel_config:
+      average_in_collective: false
+  make_sequence_length_divisible_by: 8
+  optimizer: null
+  scheduler: null
+  generation:
+    max_new_tokens: 2048
+    vllm_cfg:
+      async_engine: true
+      tensor_parallel_size: 2
+      gpu_memory_utilization: 0.5
+      expose_http_server: true
+      http_server_serving_chat_kwargs:
+        enable_auto_tools: true
+        tool_parser: hermes
+        reasoning_parser: qwen3
+    colocated:
+      enabled: false
+      resources:
+        gpus_per_node: 8
+        num_nodes: 1
+data:
+  max_input_seq_length: null
+  train:
+    data_path: "${oc.env:HF_HOME}/nanov3_data/train-split.jsonl"
+    dataset_name: NemoGymDataset
+  validation:
+    data_path: "${oc.env:HF_HOME}/nanov3_data/val-split.jsonl"
+    dataset_name: NemoGymDataset
+  default:
+    dataset_name: NemoGymDataset
+    env_name: nemo_gym
+    prompt_file: null
+    processor: nemo_gym_data_processor
+env:
+  should_use_nemo_gym: true
+  nemo_gym:
+    config_paths:
+    - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml
+    - resources_servers/math_with_judge/configs/math_with_judge.yaml
+    - resources_servers/code_gen/configs/code_gen.yaml
+    - resources_servers/workplace_assistant/configs/workplace_assistant.yaml
+    - resources_servers/mcqa/configs/mcqa.yaml
+    - resources_servers/instruction_following/configs/instruction_following.yaml
+    - resources_servers/structured_outputs/configs/structured_outputs_json.yaml
+    math_with_judge:
+      resources_servers:
+        math_with_judge:
+          judge_model_server:
+            name: policy_model
+          should_use_judge: false
+    code_gen:
+      resources_servers:
+        code_gen:
+          num_processes: 1024
+          unit_test_timeout_secs: 10
+          debug: false
+logger:
+  tensorboard_enabled: true
+  wandb:
+    project: mopd
+    name: mopd-qwen3-1.7b-3n8g-megatron-pack
+  mlflow:
+    experiment_name: mopd
+    run_name: mopd-qwen3-1.7b-3n8g-megatron-pack
+cluster:
+  gpus_per_node: 8
+  num_nodes: 3
+on_policy_distillation:
+  enabled: true
+  teacher_model_by_agent_name:
+    default_teacher: Qwen/Qwen3-1.7B
+  default_teacher_alias: default_teacher
+  strict_agent_name_match: false
+  deduplicate_shared_teacher_checkpoints: true
+  non_colocated_teachers:
+    enabled: true
+    default_teacher_cfg:
+      tensor_model_parallel_size: 2
+      pipeline_model_parallel_size: 1
+      context_parallel_size: 1
+      num_nodes: 1
+      gpus_per_node: 8
+      precision: bf16
+      micro_batch_size: 1

examples/nemo_gym/run_grpo_nemo_gym.py

@@ -208,6 +208,8 @@ def main() -> None:
         checkpointer,
         grpo_state,
         master_config,
+        teacher_worker_groups,
+        alias_to_group_alias,
     ) = setup(config, tokenizer, train_dataset, val_dataset)
 
     # NeMo-Gym is spun up inside setup() (overlapped with vLLM model load).
@@ -276,6 +278,8 @@ def main() -> None:
             grpo_save_state=grpo_state,
             master_config=master_config,
             max_trajectory_age_steps=async_config["max_trajectory_age_steps"],
+            teacher_worker_groups=teacher_worker_groups,
+            alias_to_group_alias=alias_to_group_alias,
         )
     else:
         print("🚀 Running synchronous GRPO training")

examples/run_grpo.py

@@ -142,6 +142,8 @@ def _make_policy(**kwargs):
         checkpointer,
         grpo_state,
         master_config,
+        teacher_worker_groups,
+        alias_to_group_alias,
     ) = setup(
         config,
         tokenizer,
@@ -200,6 +202,8 @@ def _make_policy(**kwargs):
             grpo_save_state=grpo_state,
             master_config=master_config,
             max_trajectory_age_steps=async_config["max_trajectory_age_steps"],
+            teacher_worker_groups=teacher_worker_groups,
+            alias_to_group_alias=alias_to_group_alias,
         )
     else:
         # Two parallel synchronous trainers (verl-style — main_ppo.py vs

examples/run_grpo_sliding_puzzle.py

@@ -262,6 +262,8 @@ def main():
         checkpointer,
         grpo_state,
         master_config,
+        _teacher_worker_groups,
+        _alias_to_group_alias,
     ) = setup(config, tokenizer, dataset, val_dataset)
 
     grpo_train(

examples/run_vlm_grpo.py

@@ -116,6 +116,8 @@ def main() -> None:
         checkpointer,
         grpo_state,
         master_config,
+        _teacher_worker_groups,
+        _alias_to_group_alias,
     ) = setup(config, tokenizer, dataset, val_dataset, processor=processor)
 
     grpo_train(