Support GEMM + Swiglu fused MLP by ksivaman · Pull Request #3971 · NVIDIA/Megatron-LM

ksivaman · 2026-03-20T19:04:31Z

What does this PR do ?

This PR supports GEMM + Swiglu fused MLP via Transformer Engine sequential ops, including:

TE GroupedTensor support in distributed optimizer.
TE GroupedTensor checkpointing support for cases to load a multi-param (1 param per expert) checkpoint into a single param config.

Note: Support is already added in dev branch and this is cherry-picking #3890.

Contribution process

Pre-checks

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

Feel free to message or comment the @mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

When your PR is ready, click Ready for Review.
An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
- Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

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.

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> Signed-off-by: ksivamani <ksivamani@nvidia.com>

copy-pr-bot · 2026-03-20T19:04:35Z

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-03-20T19:04:40Z

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

Add the oncall reviewer (optional reviewer)
Add required review teams based on your changes

See the contribution guide for more details.

Phlip79 · 2026-03-20T22:52:01Z

/claude review

Phlip79 · 2026-03-20T22:52:19Z

/ok to test ba650c7

claude

One minor typo flagged inline.

The new fused MLP path, grouped-tensor checkpoint normalization, and grouped quantized tensor handling in the distributed optimizer are substantial additions. The only test update is adding new config fields to a golden dict — there don't appear to be any unit tests covering the new logic (fused forward, _expand_quantized_param_shard_for_cast, _split_grouped_checkpoint_tensor, normalize_grouped_parameter_keys, etc.). Consider adding targeted tests for these, especially the shard-splitting and checkpoint key remapping logic which have multiple code paths.

Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com>

Phlip79 · 2026-03-24T18:35:06Z

/ok to test be3ec51

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

ksivaman · 2026-03-24T18:55:24Z

/ok to test ba650c7

copy-pr-bot · 2026-03-24T18:55:27Z

/ok to test ba650c7

@ksivaman, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

ksivaman · 2026-03-24T18:56:43Z

/ok to test d64a6bf

yashaswikarnati · 2026-03-25T04:57:01Z

    "moe_router_force_biased": None,
    "inference_grouped_gemm_backend": "auto",
    "inference_moe_disable_fused_quant_kernels": False,
+    "moe_mlp_glu_interleave_size": None,


I wonder if we should add some test when moe_mlp_glu_interleave_size is not None as this is something we are adding now ?

deepakn94 · 2026-03-25T17:18:49Z

        self.cached_grad_buffer_shard_list = [None] * len(self.buckets)
+        # Track grad mode used to create cached param views. Rebuild if mode changes to avoid
+        # mixing no_grad-created views with in-place updates in grad-enabled mode.
+        self._cached_param_buffer_shards_grad_enabled = None


What is the purpose of this?

deepakn94 · 2026-03-25T17:19:02Z

+            if self._cached_param_buffer_shards_grad_enabled != current_grad_enabled:
+                self.cached_param_buffer_shard_list = [None] * len(self.buckets)
+                self._cached_param_buffer_shards_grad_enabled = current_grad_enabled
+            with torch.no_grad():


Comment on what this is trying to do?

deepakn94 · 2026-03-25T17:19:39Z

        if self.config.use_precision_aware_optimizer_no_fp8_or_ds_fp8:
            return

+        fp8_params, shard_fp32_from_fp8, shard_offsets_in_fp8 = (


Comment on what this is trying to do?

Phlip79 · 2026-03-30T18:37:45Z

/ok to test cb7d273

Signed-off-by: Xin Yao <xiny@nvidia.com>

yaox12 · 2026-03-31T02:30:06Z

/ok to test 439688e

timmoon10 · 2026-04-01T18:00:51Z

+            else self.linear_fc2.weight0.dtype
+        )
+
+        # TODO:ksivamani: Why meta device?


@ksivaman If we didn't specify device="meta", then we would allocate memory for the weights. However, we want to reuse the weights from the GroupedLinear module.

zhongbozhu · 2026-04-22T19:58:27Z

use_transformer_engine_op_fuser is not added to the arguments.py, is that intended?

Connor-XY · 2026-05-07T04:19:03Z

Close this as we have #4636

Support GEMM + Swiglu fused MLP (NVIDIA#3890)

ba650c7

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> Signed-off-by: ksivamani <ksivamani@nvidia.com>

ksivaman requested review from a team as code owners March 20, 2026 19:04

svcnvidia-nemo-ci marked this pull request as draft March 20, 2026 19:04

ksivaman marked this pull request as ready for review March 20, 2026 19:04

svcnvidia-nemo-ci requested a review from a team March 20, 2026 19:05

svcnvidia-nemo-ci added the complexity: high label Mar 20, 2026

yaox12 approved these changes Mar 20, 2026

View reviewed changes

svcnvidia-nemo-ci added this to the Core 0.16 milestone Mar 20, 2026

claude Bot reviewed Mar 20, 2026

View reviewed changes

Comment thread megatron/core/transformer/moe/experts.py Outdated

claude Bot reviewed Mar 20, 2026

View reviewed changes

Merge branch 'NVIDIA:main' into fused_grouped_mlp_main

7b07521

erhoo82 mentioned this pull request Mar 23, 2026

[feature] enable improved Group Linear module for MoE NVIDIA-NeMo/Megatron-Bridge#2885

Closed

erhoo82 added the dev2main: mbridge dev to main: this PR is needed in main for mbridge label Mar 23, 2026

ksivaman and others added 4 commits March 23, 2026 16:24

Merge branch 'NVIDIA:main' into fused_grouped_mlp_main

661f6d9

Update megatron/core/transformer/moe/experts.py

2198c0d

Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com>

Merge branch 'main' into fused_grouped_mlp_main

619e02c

Merge branch 'main' into fused_grouped_mlp_main

be3ec51

ksivaman added 2 commits March 24, 2026 18:44

fix lint

d2b77bc

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

Merge branch 'main' into fused_grouped_mlp_main

d64a6bf

copy-pr-bot Bot temporarily deployed to test March 24, 2026 18:57 Inactive

yashaswikarnati reviewed Mar 25, 2026

View reviewed changes

deepakn94 reviewed Mar 25, 2026

View reviewed changes

Merge branch 'main' into fused_grouped_mlp_main

cb7d273

copy-pr-bot Bot temporarily deployed to test March 30, 2026 18:38 Inactive

yaox12 added 2 commits March 31, 2026 02:29

small fix

f7e17cb

Signed-off-by: Xin Yao <xiny@nvidia.com>

Merge branch 'main' into fused_grouped_mlp_main

439688e

copy-pr-bot Bot temporarily deployed to test March 31, 2026 02:31 Inactive

zhongbozhu mentioned this pull request Mar 31, 2026

[Dev] Skip routed expert padding for graph-safe MoE #4071

Merged

5 tasks

timmoon10 reviewed Apr 1, 2026

View reviewed changes

erhoo82 added the 26.04.01 label Apr 10, 2026

dingqingy-nv added core_r0.17.0 Auto-cherrypick to release branch. Apply before merge; cherrypick happens after merge. 26.06 labels Apr 10, 2026

dingqingy-nv removed the 26.04.01 label Apr 22, 2026

Mellonta requested a review from kvareddy April 24, 2026 00:46

Connor-XY mentioned this pull request May 5, 2026

Combine GEMM + SwiGLU fused MLP PRs (3890, 4071, 4095, 4219, 4311, 4324) → main #4636

Merged

3 tasks

Connor-XY closed this May 7, 2026

Uh oh!

Conversation

ksivaman commented Mar 20, 2026

What does this PR do ?

Contribution process

Pre-checks

Code review

Step 1: Mark PR as "Ready for Review"

Step 2: Final Review

Step 3: Approved

Merge

Uh oh!

copy-pr-bot Bot commented Mar 20, 2026

Uh oh!

github-actions Bot commented Mar 20, 2026

Uh oh!

Phlip79 commented Mar 20, 2026

Uh oh!

Phlip79 commented Mar 20, 2026

Uh oh!

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Phlip79 commented Mar 24, 2026

Uh oh!

ksivaman commented Mar 24, 2026

Uh oh!

copy-pr-bot Bot commented Mar 24, 2026

Uh oh!

ksivaman commented Mar 24, 2026

Uh oh!

yashaswikarnati Mar 25, 2026

Choose a reason for hiding this comment

Uh oh!

deepakn94 Mar 25, 2026

Choose a reason for hiding this comment

Uh oh!

deepakn94 Mar 25, 2026

Choose a reason for hiding this comment

Uh oh!

deepakn94 Mar 25, 2026

Choose a reason for hiding this comment

Uh oh!

Phlip79 commented Mar 30, 2026

Uh oh!

yaox12 commented Mar 31, 2026

Uh oh!

timmoon10 Apr 1, 2026

Choose a reason for hiding this comment

Uh oh!

zhongbozhu commented Apr 22, 2026

Uh oh!

Connor-XY commented May 7, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

11 participants