feat: add MOPD to single-controller text path

[yfw]

What does this PR do ?

Adds text-only Multi-Teacher On-Policy Distillation (MOPD) support to the new SingleController GRPO path.

Previously, MOPD was supported only through the legacy controller. This ports the feature to the TransferQueue-based SingleController flow while keeping teacher inference non-colocated.

Implementation

• Validates the SingleController MOPD configuration before allocating resources:
  • on_policy_distillation.enabled=true
  • grpo.adv_estimator.name=opd
  • NeMo Gym rollouts
  • non-colocated teachers
• Reserves topology-constrained policy and teacher resources before other workers can consume them.
• Creates and attaches teacher worker groups to the SingleController data plane.
• Routes each completed prompt group to its teacher using NeMo Gym agent_ref metadata.
• Runs teacher inference as a required post-write TQ enrichment stage:
  1. rollout tensors are written to TQ;
  2. teacher workers fetch their presharded rows;
  3. teachers write teacher_reference_logprobs back to the same rows;
  4. the rollout becomes available for training.
• Builds sequence-packing or dynamic-batching plans globally on the driver using logprob_mb_tokens.
• Serializes requests targeting the same physical teacher while allowing distinct teachers to run concurrently.
• Handles teacher DP padding, cleanup, cancellation, rollback, and grouped enrichment failures.
• Feeds student and teacher logprobs into the OPD advantage estimator.
• Pools OPD metrics exactly across unequal streaming chunks using sum, sum-of-squares, and token count.
• Warms the long-lived TQ partition schema before concurrent writers start.

Limitations

• This PR supports text-only MOPD. Multimodal/VLM teacher inputs are not yet forwarded through the SingleController TQ path.
• MOPD currently requires NeMo Gym and non-colocated teacher workers.
• Reference-logprob field narrowing is handled separately by fix(sc): gate reference logprobs on KL penalty #3742.

Issues

List issues that this PR closes (syntax):

Usage

• You can potentially add a usage example below

# Add a code snippet demonstrating how to use this

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

• ...

[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 5804832

[yfw]

/ok to test 7e7d721

[yfw]

Team review by 5 coordinated agents (rl-expert, bug-finder, test-agent, design-reviewer, devil's advocate) at head 7e7d721.

Overall: high-quality port — recommend addressing the two bugs before merge; nothing here is release-blocking. All 173 affected unit tests pass locally on this head, lint is green (thanks for the quick import-sort fix), recipe/nightly registration follows the conventions, and docstrings are complete.

Worth calling out as good design, explicitly: the typed PostWriteEnrichmentError enrichment seam (the buffer stays OPD-agnostic, failures classify by exception type, and the rollout retry budget is correctly not spent on post-write failures); TQTeacherLogprobCoordinator being fully unit-testable without Ray/GPUs; reuse of the legacy teacher reserve/create path with the init-ordering discipline pinned by a test; and the fail-fast MOPD validation before any resource allocation.

Findings summary: 2 bugs on config paths adjacent to (not hit by) the shipped recipe — a teacher dynamic-batching crash (executed repro, verified fix attached) and a router-replay gate that breaks MoE-student MOPD — plus 9 low-severity suggestions (1 perf-evidence ask, 2 dead-code notes, 1 stall note, 1 error-type nit, 2 design/tracking notes, 2 verified test additions). Devil's-advocate pass: 10 findings confirmed, 1 dropped as unreachable, 1 folded.

Generated by Claude Code

[yfw]

/ok to test 359c895

[terrykong]

Important

This PR now has merge conflicts with main and needs a rebase. It was MERGEABLE when this review started and went CONFLICTING while the review was being written — four commits landed on main in between, including #3771 (fix(sc): skip reference logprobs when KL is disabled), which is the most likely overlap. None of them touched opd.py or tests/unit/algorithms/test_opd.py, so the findings below still apply as written, but the line numbers will shift after you rebase.

Re-review at 359c8958 by a team of 6 coordinated agents (rl-expert, bug-finder, test-agent, design-reviewer, an architecture tracer, and a devil's advocate).

All 9 findings from the previous review are verified genuinely fixed in b0b91a3 — three lanes checked each fix independently, including sibling branches, and found no partial fixes. Nothing below re-raises them. The 5-step legacy-vs-SC convergence comparison is accepted; nothing below re-requests performance or convergence evidence.

Overall this is a well-shaped port, and the hard parts are handled carefully: post-write TQ enrichment ordering, per-physical-teacher serialization with cross-teacher concurrency, DP padding with rollback under cancellation, driver-side global packing plans that keep DP collectives in sync, and exact pooled metrics across unequal streaming chunks. 4 low-severity items follow.

Design choices we examined and think are right — please don't churn these:

• set_post_write_enricher is the right seam and the right size. A single optional callback invoked inside commit()'s existing try, so the enricher inherits the put_samples rollback path for free. We looked for a second post-write consumer and could not name one (multimodal MOPD is this same enricher widened, not a second one; reference logprobs are a train-pump stage). Turning it into a stage list would force ordering and partial-failure semantics that nothing today can pin down.
• The metric accumulators should stay inline. They score 2 of 3 on extraction — closed state and low reach-back, but no independent lifecycle. A TeacherMetrics class here would be an anemic helper. The coordinator itself earns being a class: it owns the per-teacher locks and a real stateful protocol.
• The setup ordering is pinned executably, not just in comments: assert events == ["reserve_teacher", "build_generation", "create_teacher"] turns "claim teacher GPUs before NeMo-Gym can take them, but load checkpoints after the student to avoid the shared HF-to-Megatron cache race" into a regression guard.

Two findings were raised and then killed by the adversarial pass, noted so they don't resurface: a proposed consistency guard in SingleControllerActor.__init__ (unreachable — validation enforces OPD-implies-teachers three times over) and a concern about the prev_logprobs/logprobs_policy kwarg branch (safe — validation is bidirectional and no estimator consumes both).

Caveat on verification: this review ran on macOS with no GPU, and the repo's uv.lock is linux-only, so no test and no linter was executed. Every suggestion below is marked unverified and should be run before merging.

Generated by Claude Code

[yfw]

/ok to test 37f9fed