feat(moe): Support placing MTP layers into standalone stages by BestJuly · Pull Request #2136 · NVIDIA/Megatron-LM

BestJuly · 2025-11-05T02:22:07Z

What does this PR do ?

Support placing MTP to a standalone stage, e.g., MTP in the 2nd last VPP stage for better VPP balance.
This is the PR for main branch, and the corresponding PR1916 for dev branch.

Highlights

This document outlines the design of a new MTP (Multi-Token Prediction) standalone feature. This feature allows MTP layers to be placed into a standalone VPP (Virtual Pipeline Parallelism) stage, rather than being confined to the last VPP stage.

Key Benefits:

Enhanced Flexibility in MTP VPP Placement
- MTP layers can now be positioned in either the last or second-to-last VPP stage, offering greater control over pipeline configuration.
Improved Load Balancing
- The MTP standalone feature isolates transformer block-related computations, which have a similar computational cost to a normal transformer block. This split, with MTP loss calculation remaining in the second-to-last VPP stage, contributes to a more balanced workload.
Minimal Impact on Existing Code Paths
- The design aims for near-zero impact on current code. A minor change involves the final layernorm layer, which will now be placed at the end of the decoder stage. In scenarios where the last stage contains only MTP layers and loss, this layernorm will reside in the second-to-last layer alongside a decoder layer. While this differs from previous implementations (where the final layernorm was always in the last stage), bitwise correctness has been validated under various conditions, including forceful layernorm placement and disabled gradient clipping.

Considerations:

Changes in P2P Communication Shapes
- The need to pass original and MTP hidden states to subsequent VPP stages will alter p2p communication shapes. To mitigate the impact on core communication, the shape is communicated upfront when MTP standalone settings are detected in the pipeline layout.

Design idea

We split MTP computation into two parts

Output head computation.
Others.

The advantage is that we put the MTP loss computation as main, so for the untie embedding part, the MTP loss can use exactly the same code path of the lm loss as they are in the same vpp stage, which is simpler than putting MTP loss calculation in the same vpp stage as MR#2996 so that we can avoid bugs related to this. With the new design, the DSv3 TFLOPs will also be compatible with the #2996 one which included in the 0.13 EA branch.We divide MTP computation into two parts: output head computation and others. This approach offers the advantage of treating MTP loss computation as primary. Consequently, for the untied embedding, the MTP loss can utilize the exact same code path as the LM loss, as they are within the same VPP stage. This simplifies the process compared to placing MTP loss calculation in the same VPP stage as MR#2996, thereby mitigating potential bugs. Furthermore, this new design ensures DSv3 TFLOPs compatibility with the #2996 version included in the 0.13 EA branch.

Perf and results

Perf

The throughput in DSv3 is on par or even higher than 0.13 EA (where we use previous design for this standalone feature).

Convergence

On full DSv3
(Real training mode) Almost overlapped with the main branch settings for the metrics which is related to convergence.

Deterministic correctness check
We validate the correctness in mainly two parts (when using different pp layout, can be bitwise aligned when manually disabling grad clipping)

Deterministic alignment between main branch and this MR.
- The same pipeline layout settings;
- Different pipeline layout settings.
Deterministic alignment within this MR to validate the compatibility of 1F1B overlap.

Pre-checks

I want this PR in a versioned release and have added the appropriate Milestone (e.g., Core 0.8)
I have added relevant unit tests
I have added relevant functional tests
I have added proper typing to my code Typing guidelines
I have added relevant documentation
I have run the autoformatter.sh on my PR

Code review

The following process is enforced via the CODEOWNERS file for changes into megatron/core. For changes outside of megatron/core, it is up to the PR author whether or not to tag the Final Reviewer team.

For MRs into `main` branch

(Step 1): Add PR label `Expert Review`

(Step 2): Collect the expert reviewers reviews

Attach the Expert Review label when your PR is ready for review.
GitHub auto-assigns expert reviewers based on your changes. They will get notified and pick up your PR soon.

⚠️ Only proceed to the next step once all reviewers have approved, merge-conflict are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

(Step 3): Final Review

Add Final Review label
GitHub auto-assigns final reviewers based on your changes. They will get notified and pick up your PR soon.

(Optional Step 4): Cherry-pick into release branch

If this PR also needs to be merged into core_r* release branches, after this PR has been merged, select Cherry-pick to open a new PR into the release branch.

For MRs into `dev` branch

The proposed review process for `dev` branch is under active discussion.

MRs are mergable after one approval by either eharper@nvidia.com or zijiey@nvidia.com.

Merging your PR

Any member of core-adlr and core-nemo will be able to merge your PR.

move output layer outside mtp block support mtp layer for --pipeline-model-parallel-layout run test with mtp on the second last vpp stage, with the following error message. can run, but need to check lm loss and mtp loss with baseline. fix some issue by refering to MR2996 Signed-off-by: Li Tao <lit@nvidia.com>

Signed-off-by: Li Tao <lit@nvidia.com>

…apes later Signed-off-by: Li Tao <lit@nvidia.com>

Signed-off-by: Li Tao <lit@nvidia.com>

Phlip79 · 2025-12-30T22:58:26Z

@BestJuly can you please fix conflicts?

asolergi-nv

Will take another look once conflicts are fixed & tests passing

Edit: Let's also update docs/user-guide/features/multi_token_prediction.md

asolergi-nv · 2025-12-31T09:40:03Z

+                ), "All of the MTP layers must be in the same one virtual pipeline stage"
+        for vpp_rank in range(self.virtual_pipeline_model_parallel_size - 1):
+            assert LayerType.mtp not in self.layout[0][vpp_rank], (
+                f"Corrently we restrict that the MTP should not be in the first pp rank."


Currently

Good catch! Thanks for your review and I have modified the code according to your review comments. The UT should pass and I have checked the failed one which also passed with the CI container. I believe that would be a environment issue.

============================================================================= short test summary info ============================================================================= PASSED tests/unit_tests/data/test_builder.py::test_builder

May I have your further review then? Thank you.

asolergi-nv · 2025-12-31T10:08:14Z

+
+        # print("========================")
+        # print("[DEBUG] mtp_block_spec is ", mtp_block_spec)
+        # exit()


Suggested change

# print("========================")

# print("[DEBUG] mtp_block_spec is ", mtp_block_spec)

# exit()

asolergi-nv · 2025-12-31T18:42:14Z

+    # TODO(shifangx): Currently, we only support put all of MTP layers
+    # on the last pipeline stage, so the offset is always 0.
+    # We will support more flexible MTP placement in the future.


asolergi-nv · 2025-12-31T18:44:12Z

    """Build pretraining datasets."""
    if train_valid_test_num_samples is None:
        train_valid_test_num_samples = get_train_valid_test_num_samples()
-    print_rank_0(' > datasets target sizes (minimum size):')


Why would we drop this message?

asolergi-nv · 2025-12-31T18:45:56Z

+    if vp_stage is not None:
+        return build_train_valid_test_datasets_provider(train_valid_test_num_samples, vp_stage=vp_stage)
+    else:
+        return build_train_valid_test_datasets_provider(train_valid_test_num_samples)


Do we really need this?

Signed-off-by: lit <lit@nvidia.com>

Phlip79 · 2026-01-06T21:53:08Z

/ok to test 50d34a3

asolergi-nv · 2026-01-08T12:14:22Z

/ok to test d03c64c

asolergi-nv

LGTM

shifangx and others added 14 commits November 4, 2025 18:17

Fix error when saving checkpoints

259cd21

Signed-off-by: Li Tao <lit@nvidia.com>

fix hang issue with torch.distributed.batch_isend_irecv

d515cb0

Signed-off-by: Li Tao <lit@nvidia.com>

add conditions to communicate p2p shapes, will switch to calculate sh…

ebee833

…apes later Signed-off-by: Li Tao <lit@nvidia.com>

make it compatible with 1f1b

813985d

Signed-off-by: Li Tao <lit@nvidia.com>

support mtp standalone in 1f1b

cebb164

Signed-off-by: Li Tao <lit@nvidia.com>

fix duplicated final layer norm in mtp combine

9b61229

Signed-off-by: Li Tao <lit@nvidia.com>

formatting

e650d5a

Signed-off-by: Li Tao <lit@nvidia.com>

update mtp offset in 1f1b callcables

e0dae54

Signed-off-by: Li Tao <lit@nvidia.com>

fix 1f1b overlap ut & remove final layernorm in postprocess

db1435f

Signed-off-by: Li Tao <lit@nvidia.com>

fix ut failure: no mtp_standalone attribute

29e7ea0

Signed-off-by: Li Tao <lit@nvidia.com>

add UT, fix pipeline failure and add document

8947471

Signed-off-by: Li Tao <lit@nvidia.com>

Update copyright

41e0bee

Signed-off-by: Li Tao <lit@nvidia.com>

format

c1022a7

Signed-off-by: Li Tao <lit@nvidia.com>

BestJuly requested review from a team as code owners November 5, 2025 02:22

copy-pr-bot Bot temporarily deployed to nemo-ci November 5, 2025 02:22 Inactive

ko3n1g added this to the Core 0.16 milestone Nov 5, 2025

BestJuly added module: moe Expert Review [deprecated] Apply this label to indicate that your PR is ready for expert review. labels Nov 5, 2025

copy-pr-bot Bot temporarily deployed to nemo-ci November 5, 2025 02:22 Inactive

copy-pr-bot Bot had a problem deploying to nemo-ci November 5, 2025 02:22 Failure

BestJuly assigned BestJuly and Wohox Nov 5, 2025

copy-pr-bot Bot temporarily deployed to nemo-ci November 14, 2025 03:05 Inactive

copy-pr-bot Bot temporarily deployed to nemo-ci November 14, 2025 03:36 Inactive

copy-pr-bot Bot temporarily deployed to nemo-ci November 14, 2025 03:37 Inactive

Merge branch 'main' into lit/mtp_layer_standalone_main

e3116bb

copy-pr-bot Bot temporarily deployed to nemo-ci November 18, 2025 04:20 Inactive

copy-pr-bot Bot had a problem deploying to nemo-ci November 18, 2025 04:20 Failure

copy-pr-bot Bot temporarily deployed to nemo-ci November 18, 2025 04:20 Inactive

Merge branch 'main' into lit/mtp_layer_standalone_main

5d3f862

Wohox mentioned this pull request Dec 19, 2025

[Dev] Fix ep overlap missing final layernorm #2691

Merged

6 tasks

fix missing layernorm

a3fa990

asolergi-nv reviewed Dec 31, 2025

View reviewed changes

BestJuly added 2 commits January 6, 2026 00:00

Merge branch 'main' into lit/mtp_layer_standalone_main

ae43919

Address review comments

50d34a3

Signed-off-by: lit <lit@nvidia.com>

Merge branch 'main' into lit/mtp_layer_standalone_main

69bf9bf

asolergi-nv approved these changes Jan 8, 2026

View reviewed changes

jkamalu approved these changes Jan 12, 2026

View reviewed changes

jaredcasper approved these changes Jan 14, 2026

View reviewed changes

ericharper approved these changes Jan 14, 2026

View reviewed changes

jon-barker mentioned this pull request Jan 15, 2026

Fix broken mamba-moe unit test #2970

Merged

6 tasks

xiaoyao0115 mentioned this pull request Apr 9, 2026

Minor improvements for Dynamic-cp #4226

Merged

5 tasks

yaoyu-33 mentioned this pull request May 21, 2026

[training] feat: Add Bridge support for standalone MTP pipeline stage NVIDIA-NeMo/Megatron-Bridge#3928

Merged

Navigation Menu

Uh oh!

feat(moe): Support placing MTP layers into standalone stages - #2136

feat(moe): Support placing MTP layers into standalone stages#2136
ericharper merged 22 commits into
NVIDIA:mainfrom
BestJuly:lit/mtp_layer_standalone_main

BestJuly commented Nov 5, 2025

Uh oh!

Phlip79 commented Dec 30, 2025

Uh oh!

asolergi-nv left a comment •

edited

Loading

Uh oh!

asolergi-nv Dec 31, 2025

Uh oh!

BestJuly Jan 7, 2026

Uh oh!

asolergi-nv Dec 31, 2025

Uh oh!

asolergi-nv Dec 31, 2025

Uh oh!

asolergi-nv Dec 31, 2025

Uh oh!

asolergi-nv Dec 31, 2025

Uh oh!

Phlip79 commented Jan 6, 2026

Uh oh!

asolergi-nv commented Jan 8, 2026

Uh oh!

asolergi-nv left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

11 participants

Uh oh!

Conversation

BestJuly commented Nov 5, 2025

What does this PR do ?

Highlights

Key Benefits:

Considerations:

Design idea

Perf and results

Perf

Convergence

Pre-checks

Code review

(Step 1): Add PR label Expert Review

(Step 2): Collect the expert reviewers reviews

(Step 3): Final Review

(Optional Step 4): Cherry-pick into release branch

Merging your PR

Uh oh!

Phlip79 commented Dec 30, 2025

Uh oh!

asolergi-nv left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

asolergi-nv Dec 31, 2025

Choose a reason for hiding this comment

Uh oh!

BestJuly Jan 7, 2026

Choose a reason for hiding this comment

Uh oh!

asolergi-nv Dec 31, 2025

Choose a reason for hiding this comment

Uh oh!

asolergi-nv Dec 31, 2025

Choose a reason for hiding this comment

Uh oh!

asolergi-nv Dec 31, 2025

Choose a reason for hiding this comment

Uh oh!

asolergi-nv Dec 31, 2025

Choose a reason for hiding this comment

Uh oh!

Phlip79 commented Jan 6, 2026

Uh oh!

asolergi-nv commented Jan 8, 2026

Uh oh!

asolergi-nv left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

11 participants

(Step 1): Add PR label `Expert Review`

asolergi-nv left a comment •

edited

Loading