feat: add opd/tropd metrics for super v3.5

[kbhardwaj-nvidia]

What does this PR do?

Adds proximal on-policy distillation (TROPD), comprehensive OPD diagnostics, and Super 3.5 HSG run support to GRPO.

The TROPD advantage is computed as:
log(α · π_teacher + (1 - α) · π_student) - log(π_student)

This PR:

• Adds typed and validated on_policy_distillation configuration.
• Preserves legacy OPD behavior when proximal_teacher_alpha: 1.0.
• Supports masked global-baseline subtraction.
• Adds sample-level and token-level diagnostic artifacts.
• Adds online and deferred top-k policy diagnostics, including full-vocabulary logsumexp capture.
• Adds a deterministic proxy-data builder, Super 3.5 configuration, and HSG launcher.

Usage

Configure TROPD through the top-level on_policy_distillation block:

on_policy_distillation:
  enabled: true
  use_orm_advantage: false
  proximal_teacher_alpha: 0.2
  subtract_global_baseline: true

  log_sample_stats: true
  sample_stats_log_period: 1
  log_sample_responses: false

  log_token_stats: true
  token_stats_log_period: 1

  log_topk_stats: false
  topk_stats_log_period: 5
  topk_stats_mode: student_online_teacher_deferred
  topk_stats_k: 16

  teacher_model_by_agent_name:
    example_agent: /path/to/teacher/model

Additional Information

Validation completed:

• 23 focused TROPD and diagnostics tests passed on the target branch.
• Ruff lint and format checks passed.
• Shell syntax and git diff --check passed.
• An 8-node, one-step HSG smoke completed training, refit, checkpointing, and diagnostic emission.
• A 100-step Super v24-to-Oracle TROPD run completed successfully with exit code 0:0.
• The final step-100 checkpoint was committed successfully.
• The step-100 sample artifact contains 512 valid JSONL records.
• The step-100 token artifact loads successfully with the expected diagnostic schema.

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

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[copy-pr-bot]

Auto-sync is disabled for ready for review pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[yfw]

Reviewed against base super-v3.5-posttraining (95d3cc3) at head 4dbc04c, with the merge criteria "flag-off behavior identical to base" and "new implementation correct".

Flag-off: every changed line on the default-flags path was traced and matches base — the alpha == 1.0 estimator branch (metrics get the same raw-gap tensor), _opd_seq_error_logging_fields vs the legacy masking formula (exact parity incl. the [:, 1:] slices), lm_policy.get_topk_logits torch.cat → from_batches (equal for equal-length shards; the removed "flattens" comment was already stale at base), the Megatron get_logprobs restructuring, TeacherWorkerGroup.get_topk_logits vs its sibling, and the async loop wiring. The one exception is the unconditional _opd_seq_error_logging_fields call (inline comment at grpo.py:5428), which is also what turns 15 async test_grpo.py tests red.

Correctness of the new code: TROPD math checked numerically (logaddexp form == closed form for α ∈ {1, 0.5, 0.2, 1e-3}), masked baseline finite on an all-masked batch, next-token→input-position alignment consistent with the repo's logprob convention, distributed_vocab_logsumexp matches torch.logsumexp on CPU (TP>1 path is the PR's GPU test). Remaining items are inline.

Lint: pre-commit run --all-files fails on two PR files: pyrefly at opd.py:73 (inline) and ruff-format/isort in opd_diagnostics.py:26-28 (a blank line after import torch — running pre-commit fixes it).

Smaller notes, no inline comment:

• use_orm_advantage / orm_advantage_weight (opd.py:71-72) have no readers anywhere in nemo_rl/; the only YAML setting them predates this PR and loads fine via extra="allow". Consider dropping them, or a validator rejecting non-default values, so use_orm_advantage: true cannot silently do nothing.
• The online top-k payload stores raw logits in bf16 and later subtracts the fp32 V_logsumexp, while its description calls the resulting probabilities "exact". With the shipped defer_fp32_logits: true the logits are already bf16 so nothing is lost there; for fp32-logit configs consider storing logits - lse (as the fused path does) or softening the wording.
• topk_stats_max_tokens subsamples after the full [B, S, k] top-k transfer to the driver, so it bounds the saved payload but not host memory.

Not verified here: no CUDA/Megatron/vLLM in the review environment, so the CP+packing V_logsumexp gather, the fused Megatron top-k path, and test_distributed_vocab_logsumexp were read-verified only; test_grpo.py numbers come from a scratch copy that only neutralizes the Ray-actor autouse fixture.

Generated by Claude Code