feat(mopd): add full-vocabulary on-policy distillation to SingleContr…

[RayenTian]

…oller

Adds on_policy_distillation.full, which replaces MOPD's sampled-token log-probability gap with the exact full-vocabulary reverse KL

L_t = sum_v p_student(v) * (log p_student(v) - log p_teacher(v))

This is the K=V limit of the top-k MOPD estimator: with the support spanning the whole vocabulary the score-function tail term vanishes, so the objective is exact, deterministic, and free of that estimator's off-policy bias.

Because the loss now needs the teacher's whole distribution rather than one scalar per token, the teacher/student boundary changes. Two payload paths are supported, selected by full.teacher_payload:

• hidden_states (default): the teacher ships its pre-LM-head hidden states and the student projects them with an output-layer weight shard loaded from the teacher checkpoint. The payload is hidden_size wide, so teacher and student parallelism stay fully decoupled.
• logits: the teacher ships full-vocabulary logits and the student needs no teacher LM head. The payload is vocab_size wide -- roughly 74x larger for a 2k-hidden / 152k-vocab model -- so this is a numerical-reference and fallback path, not a production configuration.

Implementation notes:

• Two chunked tensor-parallel autograd functions sit beside ChunkedDistributedEntropy and share one backward: for both weights, dw/dlog p_s is constant and cancels against the sum_v p_s = 1 normalization, leaving dL/dz = p_s * (w - L). Both recompute their log-softmaxes in backward rather than caching the teacher's, matching the neighbouring kernel.
• The reconstruction and divergence run in prepare_loss_input, the only place holding the vocab- and context-parallel groups; the loss receives a plain [B, S-1] differentiable tensor and reduces it the same way the policy-gradient path does, including the nested per-sequence average under token_level_loss=false.
• Teacher hidden states are captured with a forward pre-hook on output_layer rather than a Megatron fork. Sequence parallelism shards that input, so the capture gathers it back -- without this a TP=1 run passes while TP>1 is silently wrong.
• The payload rides a new per-token TransferQueue column, written from the pipeline stage that produced it rather than broadcast back to the replica leader, which for a per-token payload would move gigabytes for nothing.
• The teacher LM-head shard is a plain tensor on the worker, so it stays invisible to checkpoint saving and to HF/vLLM refit conversion. Its checkpoint path is resolved from the teacher's own model name: TeacherWorkerGroup copies the student policy config wholesale, so a student pretrained_checkpoint would otherwise win and the student would be distilled into itself.
• Loss knobs with no code path under this objective (ratio clipping, CISPO, importance-sampling correction, truncated IS, VAPO) are rejected at construction rather than silently ignored.

First version limits: exactly one teacher checkpoint, and the hidden-state path additionally requires student pipeline_model_parallel_size=1, because Megatron builds output_layer only on the last pipeline stage while the LM-head load is a whole-world collective -- earlier stages would fail while the last stage hangs. The logits path has no such restriction. The payload column is shaped so a per-row teacher index can be added later without reworking it.

Ships a Qwen3-1.7B self-distillation recipe whose divergence must sit at ~0, landing in disabled.txt until it has a real end-to-end run and the nightly GPU-hour budget has room for it.

What does this PR do ?

Add a one line overview of what this PR aims to accomplish.

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.