feat: add MTP (multi-token prediction) training support for nemotron by yfw · Pull Request #2801 · NVIDIA-NeMo/RL

yfw · 2026-06-12T19:45:29Z

What does this PR do ?

Adds training-side Megatron Multi-Token Prediction support, rebased onto the Megatron-Bridge 79c8957 bump which includes needed Megatron-LM changes. Part of Nemotron 3 Super upstreaming.

megatron_policy_worker: compute mtp_loss_mask from token_mask × sample_mask (with optional mtp_positive_only gating on advantages), set mtp_grad_scale_func for MTP loss scaling, and collect MTP metrics.
data.py / train.py: plumb mtp_loss_mask through microbatch processing/packing into model_forward's loss_mask.
common.py: get_mtp_metrics (per-layer MTP loss + acceptance rate).
grpo.py / lm_policy.py: surface mtp/* metrics in training logs.

Scope: training only. The MTP inference path (vLLM drafter weight load / speculative-config handling) will be handled separately as part of the ultra upstreaming.

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 · 2026-06-12T19:45:33Z

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 · 2026-06-12T19:46:31Z

✅ Submodule Fast-Forward Check Results

Check based on commit: 05c85eb (PR #2801 from yifu/super_mtp_merge_main)

✅ Submodules that are properly updated:

Megatron-Bridge: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

yfw · 2026-06-12T19:47:38Z

/ok to test 05c85eb

yfw · 2026-06-13T06:06:25Z

/ok to test 2fc8700

github-actions · 2026-06-13T06:06:49Z

✅ Submodule Fast-Forward Check Results

Check based on commit: 2fc8700 (PR #2801 from yifu/super_mtp_merge_main)

✅ Submodules that are properly updated:

Megatron-Bridge: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

Cherry-pick of 8106050 (MTP changes) onto ci/bump-mbridge-79c8957. Resolved conflicts against the newer base; build-custom-vllm.sh change dropped per request. Where the commit's context lines were stale (master_config access style, model_forward cfg= arg, vllm hf_overrides.update, the inline IPC weight loader), kept the new base's versions and grafted only the genuine MTP additions (mtp_loss_mask plumbing, mtp_metrics logging, num_speculative_tokens==0 handling, megatron parallel-state imports). Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com> Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

- test_mtp_metrics.py: get_mtp_metrics per-layer loss/acceptance (1-indexed), empty-tracker and defaults-when-only-loss cases. Stubs the cross-rank reduce so it runs single-process on CPU. - test_megatron_data.py: process_microbatch propagates mtp_loss_mask in the no-packing path and packs it alongside input_ids in the packing path; absent -> None. Tightened the existing with_packing fixture's __contains__ mock (was return_value=True for all keys, which now spuriously triggers the mtp_loss_mask branch). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

setup._apply_mtp_config currently only wires mtp_num_layers; extend it to also apply mtp_loss_scaling_factor, mtp_use_repeated_layer and mtp_detach_heads onto the mcore model config (gated on presence), so these MTP knobs are user-controllable instead of silently using megatron defaults. Ported from 499108c. Config conventions: - Declare the MTP keys (incl. mtp_positive_only) on the MegatronConfig TypedDict and add documented defaults to the exemplar examples/configs/grpo_math_1B.yaml (+ matching tests/unit/reference_configs/grpo_math_1B.yaml so the v1->v2 check stays green). - Drop the forbidden call-site default in megatron_policy_worker: cfg["megatron_cfg"].get("mtp_positive_only", False) -> .get("mtp_positive_only"). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

mtp_positive_only (restrict the MTP loss mask to positive-advantage tokens) was experimental and is unused. Remove it from the MTP loss-mask computation, the MegatronConfig schema, and the grpo_math_1B exemplar/reference configs. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

Collecting MTP metrics inline in MegatronPolicyWorker.train() made cloudpickle pull an unpicklable torch ConfigModuleInstance into the actor serialization (`cannot pickle 'ConfigModuleInstance' object`), so Ray could not create the megatron policy worker. Move the collection into a small _collect_mtp_metrics method that lazily imports get_mtp_metrics, and drop the module-level get_mtp_metrics import, so train()'s (large) globals footprint no longer references it. Verified the worker class cloudpickles again. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

ashors1

Thanks for the MTP training support — clean implementation with good test coverage!

A few findings below (5 inline comments). The VLM guard and test assertion are the most actionable.

Copyright note (spans multiple files): common.py, data.py, megatron_policy_worker.py, __init__.py, lm_policy.py, and grpo.py still carry # Copyright (c) 2025. Current year is 2026 — train.py and setup.py were already correctly updated in this PR.

Generated by Claude Code

Co-authored-by: Anna Shors <ashors@nvidia.com> Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

yfw · 2026-06-15T23:10:33Z

/ok to test 0981ab0

Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

yfw · 2026-06-16T17:57:16Z

/ok to test b86c1a1

Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Co-authored-by: Anna Shors <ashors@nvidia.com> Signed-off-by: Anna Shors <ashors@nvidia.com>

yfw requested review from a team as code owners June 12, 2026 19:45

yfw added the CI:L1 Run doctests, unit tests, and functional tests label Jun 12, 2026

copy-pr-bot Bot temporarily deployed to public June 12, 2026 19:47 Inactive

copy-pr-bot Bot temporarily deployed to public June 12, 2026 19:48 Inactive

copy-pr-bot Bot temporarily deployed to test June 12, 2026 19:51 Inactive

copy-pr-bot Bot temporarily deployed to public June 12, 2026 19:51 Inactive

yfw marked this pull request as draft June 12, 2026 19:58

yfw marked this pull request as ready for review June 13, 2026 06:06

copy-pr-bot Bot temporarily deployed to public June 13, 2026 06:06 Inactive

copy-pr-bot Bot temporarily deployed to public June 13, 2026 06:07 Inactive

copy-pr-bot Bot temporarily deployed to test June 13, 2026 06:09 Inactive

copy-pr-bot Bot temporarily deployed to public June 13, 2026 06:10 Inactive

yfw and others added 5 commits June 15, 2026 11:09

yfw force-pushed the yifu/super_mtp_merge_main branch from 2fc8700 to f22fece Compare June 15, 2026 18:23

yfw removed the CI:L1 Run doctests, unit tests, and functional tests label Jun 15, 2026

copy-pr-bot Bot temporarily deployed to public June 15, 2026 18:33 Inactive

copy-pr-bot Bot temporarily deployed to public June 15, 2026 18:37 Inactive

copy-pr-bot Bot temporarily deployed to test June 15, 2026 18:37 Inactive

ashors1 reviewed Jun 15, 2026

View reviewed changes

Comment thread nemo_rl/models/megatron/common.py Outdated

Comment thread nemo_rl/models/megatron/data.py

Comment thread nemo_rl/models/megatron/data.py

Comment thread tests/unit/models/megatron/test_megatron_data.py Outdated

Comment thread nemo_rl/models/megatron/common.py

macandro96 reviewed Jun 15, 2026

View reviewed changes

Comment thread nemo_rl/models/policy/workers/megatron_policy_worker.py

yfw and others added 2 commits June 15, 2026 15:59

Update nemo_rl/models/megatron/common.py

bae7916

Co-authored-by: Anna Shors <ashors@nvidia.com> Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

Address review comments

0981ab0

Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

copy-pr-bot Bot temporarily deployed to public June 15, 2026 23:10 Inactive

copy-pr-bot Bot temporarily deployed to public June 15, 2026 23:11 Inactive

copy-pr-bot Bot temporarily deployed to test June 15, 2026 23:14 Inactive

copy-pr-bot Bot temporarily deployed to public June 15, 2026 23:15 Inactive

Merge branch 'main' into yifu/super_mtp_merge_main

b86c1a1

Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

macandro96 approved these changes Jun 16, 2026

View reviewed changes

copy-pr-bot Bot temporarily deployed to public June 16, 2026 18:02 Inactive

copy-pr-bot Bot temporarily deployed to public June 16, 2026 18:03 Inactive

copy-pr-bot Bot temporarily deployed to test June 16, 2026 18:05 Inactive

copy-pr-bot Bot temporarily deployed to public June 16, 2026 18:07 Inactive

ashors1 approved these changes Jun 16, 2026

View reviewed changes

yfw changed the title ~~feat: add MTP (multi-token prediction) training support~~ feat: add MTP (multi-token prediction) training support for nemotron Jun 16, 2026

yfw merged commit d584a89 into main Jun 16, 2026
76 of 77 checks passed

yfw deleted the yifu/super_mtp_merge_main branch June 16, 2026 21:11

sharonyu-115 mentioned this pull request Jun 18, 2026

[bug] MTP loss mask trips "MTP not supported with VLM sequence packing" assertion on non-MTP Qwen3.5 megatron recipes #2869

Closed

yfw mentioned this pull request Aug 27, 2026

feat(megatron): support MTP in split training #3881

Merged

4 tasks

Navigation Menu

Uh oh!

feat: add MTP (multi-token prediction) training support for nemotron - #2801

feat: add MTP (multi-token prediction) training support for nemotron#2801
yfw merged 8 commits into
mainfrom
yifu/super_mtp_merge_main

yfw commented Jun 12, 2026

Uh oh!

copy-pr-bot Bot commented Jun 12, 2026

Uh oh!

github-actions Bot commented Jun 12, 2026

Uh oh!

yfw commented Jun 12, 2026

Uh oh!

yfw commented Jun 13, 2026

Uh oh!

github-actions Bot commented Jun 13, 2026

Uh oh!

ashors1 left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

yfw commented Jun 15, 2026

Uh oh!

yfw commented Jun 16, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

yfw commented Jun 12, 2026

What does this PR do ?

Issues

Usage

Before your PR is "Ready for review"

Additional Information

Uh oh!

copy-pr-bot Bot commented Jun 12, 2026

Uh oh!

github-actions Bot commented Jun 12, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

yfw commented Jun 12, 2026

Uh oh!

yfw commented Jun 13, 2026

Uh oh!

github-actions Bot commented Jun 13, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

ashors1 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

yfw commented Jun 15, 2026

Uh oh!

yfw commented Jun 16, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants