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

[yfw]

What does this PR do ?

Adds Multi-Teacher On-Policy Distillation (MOPD) on top of async GRPO. MOPD distills one or more teacher models into the policy using token-level teacher-vs-student logprob advantages, with teachers served by non-colocated inference worker groups. Teacher routing is per-agent (via NeMo-Gym), so different environments/agents can be distilled from different teachers.

Algorithm

• OPDAdvantageEstimator (nemo_rl/algorithms/advantage_estimator.py): token-level distillation advantage Â_t = sg[log π_teacher − log π_student]. Importance-sampling truncation is handled separately by the ICE-POP gate in the loss function.
• The MOPD update is expressed through loss_fn: disable_ppo_ratio (REINFORCE form) + use_importance_sampling_correction with truncated_importance_sampling_type: icepop. When adv_estimator.name: opd is set without these, _create_advantage_estimator emits warnings.

Teacher infrastructure

• nemo_rl/algorithms/opd.py (new): on_policy_distillation config schemas (pydantic BaseModel), teacher routing (agent → teacher-alias resolution and deduplication of shared checkpoints), teacher worker-group creation, and the TP-aligned sequence-pad helper for non-packed teachers.
• nemo_rl/models/policy/teacher_worker_group.py (new): TeacherWorkerGroup, an inference-only Megatron worker group (no optimizer, no reference model) that loads a teacher checkpoint and exposes get_logprobs, plus config builders with per-teacher overrides and dedup.
• nemo_rl/models/megatron/setup.py: freezes all parameters for inference-only workers so DDP skips gradient-buffer allocation (significant memory savings for large models).
• Note: MOPD is gym-coupled by design (per-agent teacher routing enables multi-teacher / multi-environment distillation).

Integration

• on_policy_distillation is now a typed field on the GRPO MasterConfig.
• setup() reserves nodes for non-colocated teachers and returns the teacher worker groups plus the alias→group mapping.
• The async trajectory collector computes teacher logprobs at collection time (overlapped with rollouts), with per-teacher locks, data-parallel padding, and TP-aligned sequence padding; teacher logprobs are plumbed through to the advantage computation and padded to the training sequence length.

Config, examples, and tests

• Documented the on_policy_distillation field in the grpo_math_1B exemplar (and reference config).
• New recipe examples/configs/recipes/llm/mopd-qwen3-1.7b-3n8g-megatron-pack.yaml: dense Qwen3-1.7B self-distill (student == teacher → OPD loss ≈ 0) on 3 nodes (1 policy + 1 vLLM + 1 teacher), added to the nightly suite. The driver asserts the loss stays near zero. mopd is registered as a recipe algo prefix.
• Unit tests for the OPD advantage estimator, OPD helpers/routing/config, the teacher worker-group config builders, and collection-time teacher-logprob padding/routing.

Issues

List issues that this PR closes (syntax):

Usage

MOPD requires async GRPO with non-colocated teachers (NeMo-Gym). Enable it via:

grpo:
  async_grpo: { enabled: true }
  adv_estimator: { name: opd }
loss_fn:
  disable_ppo_ratio: true
  use_importance_sampling_correction: true
  truncated_importance_sampling_type: icepop
on_policy_distillation:
  enabled: true
  teacher_model_by_agent_name:
    default_teacher: <teacher-checkpoint>
  default_teacher_alias: default_teacher
  non_colocated_teachers:
    enabled: true
    default_teacher_cfg: { tensor_model_parallel_size: 2, num_nodes: 1 }

See examples/configs/recipes/llm/mopd-qwen3-1.7b-3n8g-megatron-pack.yaml for a complete example.

Additional Information

• Reference: MOPD
• Example reproduction of the MOPD phase in nvidia/Nemotron-Cascade-2-30B-A3B: https://github.com/NVIDIA-NeMo/RL/blob/e77dcdcabc71181031a7cb1e2518e224a195fc63/docs/guides/mopd-cascade.md

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[yfw]

/ok to test 99b0186

[yfw]

/ok to test 7203c79

[yfw]

/ok to test 8d0844c

[yfw]

/ok to test ff604ae

[qiaochuz-nv]

RL POR Test Summary — MOPD (Multi-Teacher On-Policy Distillation)

Feature under test: NVIDIA-NeMo/RL PR #2780 — Multi-Teacher On-Policy Distillation
Merge commit pinned in driver: 5c2322cfc4adc26538ebd976f5ce853dc7733a51
Test case: test_grpo_mopd_qwen3_1p7b_megatron_multinode
MR: https://gitlab-master.nvidia.com/dl/DLQASH/nmfw_tests/-/merge_requests/267 (branch qiaochuz/por_2780_mopd, commit 0ec559f)

---

Result: PASS ✅

Field	Value
Status	passed
Exit code	0
Failures	0
Run ID	20260629T085457Z-355b556f
SLURM job	5528760
Layout	3 nodes / 24 GPUs (batch, exclusive, Ray bring-up)
Duration	997 s (~16.6 min)
Container	nemo-rl-nightly-20260629.sqsh (commit 60337ec), built AFTER PR #2780 merge
Cluster	EOS (login-eos02), account coreai_dlalgo_ci
Start / End (UTC)	2026-06-29T09:21:07Z → 09:37:44Z

Source: logs/20260629T085457Z-355b556f/tests/test_grpo_mopd_qwen3_1p7b_megatron_multinode/summary.json

---

What the test exercises

Drives the shipped recipe directly — no fork of product config:
examples/configs/recipes/llm/mopd-qwen3-1.7b-3n8g-megatron-pack.yaml
via examples/nemo_gym/run_grpo_nemo_gym.py at the native 3-node / 24-GPU topology:

• 1 node — Megatron policy student (Qwen/Qwen3-1.7B)
• 1 node — vLLM generation
• 1 node — non-colocated bf16 Megatron teacher worker group (TP=2, PP=1, CP=1)

Self-distillation (student == teacher == Qwen3-1.7B), so the OPD advantage loss
is ≈ 0 by construction → a correctness smoke test for the new code paths, not
a convergence test.

MOPD code paths covered

• adv_estimator.name = opd (on-policy distillation advantage)
• Async GRPO (async_grpo.enabled = true) — async trajectory collector + replay buffer
• Non-colocated teacher worker group (non_colocated_teachers.enabled = true)
• Teacher logprob path (teacher_model_by_agent_name.default_teacher)
• Sequence-level logprob error masking (seq_logprob_error_threshold = 2.0)
• ICEPOP truncated importance sampling, reference-policy KL penalty = 0.0

---

Real per-step evidence (stdout.log)

▶ Applying sequence-level logprob error masking (threshold=2.0)...
  Masked 0 sequences with mult_prob_error > 2.0
[teacher_logprob] group=default_teacher samples=2 ... across 1 teacher(s)
📊 Advantages stats: min=0.0000, max=0.0000, mean=0.0000, std=0.0000   # self-distill → 0

Step 1 — Training Results:
  • Loss: 0.0000        • Generation KL Error: 0.0018
  • Avg Reward: 0.2500  • Buffer Size: 4

Step 2 — Training Results:
  • Loss: 0.0000        • Generation KL Error: 0.0019
  • Avg Reward: 0.0000  • Buffer Size: 4

• OPD adv estimator active, advantages exactly 0 (expected for self-distill)
• Teacher logprob computed via the non-colocated teacher (default_teacher)
• Max Generation KL Error ~0.0018–0.0019 (well under masking threshold 2.0)
• In-script validators all PASS (missing=[], 2 training blocks, losses [0.0, 0.0])

---

CI shrink (product behavior untouched)

Only step count and sequence length reduced to fit CI:

• grpo.max_num_steps reduced (smoke: 2 steps)
• policy.max_total_sequence_length = 8192 (floor — 4096 hangs), with
  vllm_cfg.max_model_len and mcore_generation_config.max_model_len
  set to the same 8192 to keep generation backends consistent.

MOPD knobs left at recipe defaults: adv_estimator, async_grpo, teacher
config, ICEPOP, KL penalty, masking threshold — all as shipped.

---

Notes / observations

1. group_raw_exit_code=1 in notes is a benign Ray finalize rc (cluster
   teardown), not a test failure — harness reports status=passed failures=0.
   The driver's log validator keys off training markers in run.log, not the
   Ray process exit code.
2. Doc nit found during verification (filed separately as RL doc-fix PR): the
   docs/about/algorithms/mopd.md CLI example uses data.val.data_path, but the
   shipped recipe key is data.validation.data_path — verbatim copy-paste would
   hit an OmegaConf "key not found".

---

Registration

"test_grpo_mopd_qwen3_1p7b_megatron_multinode": {
  "tier": "e2e", "submit_mode": "sbatch", "num_gpus": 24, "timeout_sec": 9000,
  "slurm": { "nodes": 3, "exclusive": true, "ray_bringup": true },
  "tags": ["codecoverage", "rc_p0"],
  "status": "new", "version": "nightly",
  "pr": "https://github.com/NVIDIA-NeMo/RL/pull/2780"
}