feat(megatron): support MTP in split training by yfw · Pull Request #3881 · NVIDIA-NeMo/RL

yfw · 2026-08-27T19:52:17Z

What does this PR do ?

Ports the Megatron multi-token prediction (MTP) training lifecycle from #2801 into the SingleController split training path.

Derives the MTP loss mask from the existing token and sample masks before microbatch packing.
Leaves mtp_grad_scale_func unset so MCore inherits the optimizer's loss scale, including fp16 dynamic scaling.
Normalizes detached MTP auxiliary gradients by the global valid-token count independently of the main loss denominator.
Rejects nonzero attached-head MTP with a sequence-level main loss because that mixed backbone gradient cannot be corrected after split-step accumulation.
Collects the independently clipped MTP gradient norm plus per-layer loss and acceptance metrics after the optimizer step.
Propagates worker-reduced MTP metrics through TQPolicy to the training logger.
Forwards model-owned sequence/MTP-mask packing and context-parallel slicing flags through split training.

No TQ schema or new user configuration fields are required.

Issues

Follow-up to #2801.

Split path: support MTP + SEQUENCE_LEVEL loss with mtp_detach_heads=False #3889 tracks support for attached-head MTP with a sequence-level main loss.
Split path: MoE auxiliary loss inherits the main loss denominator instead of 1/global_valid_toks #3888 tracks the analogous pre-existing MoE auxiliary-loss denominator issue.

Usage

Existing MTP-enabled Megatron recipes now use the same MTP training behavior when driven through the SingleController split API.

Before your PR is "Ready for review"

Pre checks:

Read and followed the contributor guidelines.
Added focused lifecycle and TQ propagation tests.
Ran the available focused unit tests and static checks locally.
No documentation update is needed because this restores parity for existing configuration.

Additional Information

Validation:

Focused unit tests: 11 passed, 1 skipped.
Ruff lint and formatting checks passed.
Repository pre-commit hooks, including Pyrefly, passed.
Python compilation and git diff checks passed.

The MCore-specific state-machine tests were skipped locally because Transformer Engine is unavailable on this host. Installing the full MCore extras is also blocked here by the missing infiniband/mlx5dv.h system header.

Sample run with single controller and mtp training enabled:

wandb: https://wandb.ai/nvidia/nano35-rlvr-convergence/runs/nk99s0qe

Add MTP loss masking and gradient scaling to the split Megatron train-step lifecycle. Collect MTP losses, acceptance rates, and gradient norms at finish, then propagate them through TQPolicy. Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

copy-pr-bot · 2026-08-27T19:52:21Z

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

The MTP lifecycle port looks sound. scale_gradients does reach the MTP parameters (it iterates buffers + expert_parallel_buffers and scales each whole buffer, with no MTP or expert exclusion), so the raw-accumulate-per-chunk then normalize-once-at-finish contract holds. The loss mask is load-bearing too — without it process_mtp_loss gets loss_mask=None and trains MTP over prompt and padding tokens.

Three inline comments: a latent fp16 issue in the gradient-scale hook, a test assertion that cannot currently fail, and one intent question about auxiliary-branch normalization that is pre-existing on this path.

Generated by Claude Code

Let MCore inherit the optimizer loss scale for MTP and normalize detached MTP gradients by valid tokens when the main loss uses a sequence denominator. Reject the nonzero attached-head combination that cannot be corrected after split-step accumulation. Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

yfw · 2026-08-28T00:04:09Z

/ok to test 60c55c8

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

yfw · 2026-09-01T17:10:26Z

/ok to test 3ea9fd7

mehraakash

The core MTP lifecycle and scaling logic looks consistent with the synchronous path: mask construction, fp16 loss-scale inheritance, detached-head token normalization, metric reduction, and TQ propagation all look sound.

I left one compatibility question about forwarding the model-owned packing and context-parallel flags through the split microbatch path.

Could we also run CI:L1 against the exact current head and include a real MCore split-MTP functional test? The current head has only metadata, security, and submodule checks. The previous Lfast PPO failure used a non-MTP configuration and appears unrelated, but there is not yet a clean functional signal for this training-path change.

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

Keep the SingleController microbatch iterator aligned with synchronous training for models that own sequence packing, MTP-mask packing, or context-parallel slicing. Add regression coverage for both supported capability combinations. Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

yfw · 2026-09-03T04:26:27Z

/ok to test b93b84c

yfw · 2026-09-03T04:26:38Z

I triggered CI:L1 for the exact updated head, b93b84c. I did not add a nominal MTP L1 script because the current SingleController L1 recipe uses Qwen3-0.6B with mtp_num_layers=0 and no MTP weights, so it would not exercise the behavior under review. The repository’s actual Megatron MTP functional recipe requires the private Super Omni checkpoint/training data and 16 nodes, and is currently disabled. The added MCore regression instead covers propagation for both supported model-owned packing/CP-slicing combinations; a true end-to-end CI lane will require a small CI-accessible MTP checkpoint.

yfw · 2026-09-03T04:28:13Z

/ok to test b93b84c

yfw · 2026-09-03T04:28:45Z

/ok to test b93b84c

yfw · 2026-09-03T04:29:32Z

CI:L1 is explicitly dispatched against b93b84c: https://github.com/NVIDIA-NeMo/RL/actions/runs/33715206186

…t tests Two unit tests drifted from code this branch added and went red in CI (run 33939587768 on 30d300c): - test_train_pump_e2e::_prepopulate_buffer inserts a ready slot into TQReplayBuffer by appending to its parallel lists by hand. The branch added _rollout_ids_list and _staging_keys_list to the buffer, so the first remove() indexed past the end of the empty staging list (IndexError in _remove_unlocked). Append None to both, matching what begin_group/commit and checkpoint restore do for finalized groups. - test_rollout_pump_writes_expected_tq_data pins the TQ tag schema. The branch added num_routed_experts_backfilled to VIOLATION_TAG_KEYS so every row's tags now carry it; add it to the expected set. The remaining red shard, L0_Unit_Tests_Megatron_Policy_1, is a main-side conflict between #3881 (test expects train_microbatch to forward model_slices_context_parallel_inputs) and #2957 (guard raises on it) and is inherited unchanged. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Pranav Thombre <pthombre@nvidia.com>

feat(megatron): support MTP in split training

b0dd481

Add MTP loss masking and gradient scaling to the split Megatron train-step lifecycle. Collect MTP losses, acceptance rates, and gradient norms at finish, then propagate them through TQPolicy. Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

yfw requested review from a team as code owners August 27, 2026 19:52

yfw commented Aug 27, 2026

View reviewed changes

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

Comment thread tests/unit/models/policy/test_megatron_split_state.py Outdated

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

This was referenced Aug 27, 2026

Split path: MoE auxiliary loss inherits the main loss denominator instead of 1/global_valid_toks #3888

Open

Split path: support MTP + SEQUENCE_LEVEL loss with mtp_detach_heads=False #3889

Open

yfw added the CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) label Aug 28, 2026

Merge remote-tracking branch 'origin/main' into yifu/v2_mtp_training

3ea9fd7

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

Merge branch 'main' into yifu/v2_mtp_training

babac99

mehraakash reviewed Sep 2, 2026

View reviewed changes

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

yfw added 2 commits September 2, 2026 21:21

Merge remote-tracking branch 'origin/main' into yifu/v2_mtp_training

86602ac

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

yfw removed the CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) label Sep 3, 2026

yfw added the CI:L1 Run doctests, unit tests, and functional tests label Sep 3, 2026

mehraakash approved these changes Sep 4, 2026

View reviewed changes

yfw merged commit 533d08e into main Sep 4, 2026
110 of 116 checks passed

yfw deleted the yifu/v2_mtp_training branch September 4, 2026 18:59

terrykong mentioned this pull request Sep 5, 2026

refactor(docker): read the actor venv list from a Python leaf module #4002

Merged

terrykong mentioned this pull request Sep 7, 2026

ci: Bump Megatron-Bridge to 5ed9799 #4022

Merged

4 tasks

Navigation Menu

Uh oh!

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

feat(megatron): support MTP in split training#3881
yfw merged 6 commits into
mainfrom
yifu/v2_mtp_training

yfw commented Aug 27, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Aug 27, 2026

Uh oh!

yfw left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

yfw commented Aug 28, 2026

Uh oh!

yfw commented Sep 1, 2026

Uh oh!

mehraakash left a comment

Uh oh!

Uh oh!

yfw commented Sep 3, 2026

Uh oh!

yfw commented Sep 3, 2026

Uh oh!

yfw commented Sep 3, 2026

Uh oh!

yfw commented Sep 3, 2026

Uh oh!

yfw commented Sep 3, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

yfw commented Aug 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do ?

Issues

Usage

Before your PR is "Ready for review"

Additional Information

Uh oh!

copy-pr-bot Bot commented Aug 27, 2026

Uh oh!

yfw left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

yfw commented Aug 28, 2026

Uh oh!

yfw commented Sep 1, 2026

Uh oh!

mehraakash left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

yfw commented Sep 3, 2026

Uh oh!

yfw commented Sep 3, 2026

Uh oh!

yfw commented Sep 3, 2026

Uh oh!

yfw commented Sep 3, 2026

Uh oh!

yfw commented Sep 3, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

yfw commented Aug 27, 2026 •

edited

Loading