Skip to content

feat(mopd): add cross-tokenizer teacher scoring - #4159

Draft
avenkateshha wants to merge 4 commits into
NVIDIA-NeMo:avenkateshha/xtoken-offset-alignmentfrom
avenkateshha:avenkateshha/mopd-xtoken-merge
Draft

avenkateshha wants to merge 4 commits into
NVIDIA-NeMo:avenkateshha/xtoken-offset-alignmentfrom
avenkateshha:avenkateshha/mopd-xtoken-merge

Conversation

@avenkateshha

@avenkateshha avenkateshha commented Sep 16, 2026 •

Copy link
Copy Markdown
Contributor

What does this PR do?

Adds cross-tokenizer MOPD teacher scoring on top of #3286's offset aligner.

Important

This is a true stacked draft PR. Its base is avenkateshha/xtoken-offset-alignment at ca3366fc5cf99add7cc4c2002346946f83a5ee85; the diff contains only MOPD work. After #3286 merges, this same branch/PR will be rebased and retargeted to refreshed main.

The implementation:

  • reconstructs the exact sampled student transcript and renders a teacher-native transcript;
  • aligns generated assistant turns only, then sum-preservingly projects teacher log-probabilities onto student positions;
  • keeps sampled IDs authoritative from the runtime tokenizer while loading independent plain-Hugging-Face student and teacher alignment tokenizers, including when Fastokens is active;
  • integrates with the current NemoGymRolloutResult.final_batch collector contract and stores both projected scores and a boolean validity mask in replay;
  • intersects the teacher-validity mask only inside OPD advantages, without changing the policy loss mask or denominator;
  • preserves same-token and full-vocabulary MOPD behavior, strict teacher routing/deduplication, and replay compatibility;
  • adds Mingyu's optional TOP-D scalar reward shaping as a separate default-off commit (this is not full TOP-D);
  • adds focused public tests, observability, documentation, and a placeholder-based manual/L2-only parity recipe.

PR #4045 remains an integration/sequencing constraint only. This PR does not cherry-pick or manage it, and preserves its same-token/full-vocabulary stable-index behavior for a later rebase if needed.

Issues

Depends on #3286.

Usage

Replace the public placeholders in:

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

Then launch through the existing NeMo Gym GRPO entrypoint. Cross-tokenizer MOPD intentionally rejects Single-Controller, full-vocabulary, and multimodal combinations in this first version.

Validation

Tested through 6d95861649d4d58ad5a6dbc889d5b7ac2a771e6d:

  • 95 passed: focused X-Token scorer/helpers, OPD advantages, routing/configuration, and teacher-worker tests.
  • 8 passed: replay score/mask compatibility and current NemoGymRolloutResult.final_batch success/fail-closed cases.
  • 56 passed, 1 deselected: existing TokenAligner and same-token full-vocabulary OPD suites. The deselected case reaches a pre-existing missing NATIVE_MULTIMODAL_KEYS import, reproduced on the stack base.
  • 704 passed, 6 skipped: full config validation and config-v2 suites.
  • Recipe/suite accounting and naming passed; the new manual driver passes TEST_DRYRUN=1.
  • Ruff check and format check passed (718 files already formatted).
  • Pyrefly reaches only the same pre-existing stack-base error: rollout_manager.py imports missing NATIVE_MULTIMODAL_KEYS from multimodal_utils.

Internal CPU differential parity used all 192 corrected private records without committing any record, prompt, or tool trace. Runtime Fastokens plus separate plain-HF alignment tokenizers produced 192/192 exact transcript/teacher IDs, offset shapes, #3286-vs-16e4cc alignment coordinates, pairs/correctness, template masks, final validity masks, and recorded-logprob projections at 1e-6 tolerance (maximum absolute/relative projection error 0).

Pending explicit command/allocation approval: BF16 teacher-forward parity, frozen replay advantage/loss parity, GPU teacher smoke plus checkpoint/resume, and the three seed-42 multi-node curve runs. No costly GPU or multi-node job was launched.

Before this PR is ready for review

  • Read and followed the contributor guidelines.
  • Added focused tests and reused the existing aligner/full-vocabulary suites.
  • Ran the relevant CPU unit, config, lint, and type-check gates locally.
  • Updated the MOPD documentation and added a public placeholder recipe.
  • Run the explicitly approved GPU/L2 qualification listed above.

Additional information

Signed-off-by: Adithya Hanasoge <avenkateshha@nvidia.com>
Signed-off-by: Adithya Hanasoge <avenkateshha@nvidia.com>
Signed-off-by: Adithya Hanasoge <avenkateshha@nvidia.com>
@copy-pr-bot

copy-pr-bot Bot commented Sep 16, 2026

Copy link
Copy Markdown

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.

@github-actions github-actions Bot added the Documentation Improvements or additions to documentation label Sep 16, 2026
@avenkateshha avenkateshha added the CI:L0 Run doctests and unit tests label Sep 16, 2026
@avenkateshha

Copy link
Copy Markdown
Contributor Author

/ok to test 71cad59

Signed-off-by: Adithya Hanasoge <avenkateshha@nvidia.com>
@avenkateshha

Copy link
Copy Markdown
Contributor Author

/ok to test 6d95861

@avenkateshha

Copy link
Copy Markdown
Contributor Author

CI update for 6d95861649d4d58ad5a6dbc889d5b7ac2a771e6d: the rerun now passes recipe minimization, Ruff/import formatting, Ruff format, TOML formatting, and Markdown checks. The sole lint error is the pre-existing stack-base Pyrefly failure in nemo_rl/experience/rollout_manager.py:42 (NATIVE_MULTIMODAL_KEYS is absent from multimodal_utils), which I reproduced before the port at ca3366fc5cf99add7cc4c2002346946f83a5ee85. The L0 quality gate propagates that baseline failure; no L2/convergence job was launched.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CI:L0 Run doctests and unit tests Documentation Improvements or additions to documentation

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant