Skip to content

[Main] Support device-init grouped linear module too with when not using TE opfuser - #6000

Open
zhongbozhu wants to merge 7 commits into
NVIDIA:mainfrom
zhongbozhu:main_support_device_group_linear
Open

[Main] Support device-init grouped linear module too with when not using TE opfuser#6000
zhongbozhu wants to merge 7 commits into
NVIDIA:mainfrom
zhongbozhu:main_support_device_group_linear

Conversation

@zhongbozhu

@zhongbozhu zhongbozhu commented Jul 23, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor
  • I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Proper single weight support, discrete weight is not affected. Adds many numerical unit tests running on GB200.

Unit testing:

Note: some tests can only run with TE updated (NVIDIA/TransformerEngine#3224), for CI image, it should just pass without causing errors.

pytest -s -v --confcutdir=tests/unit_tests/transformer/moe tests/unit_tests/transformer/moe/test_grouped_mlp.py

torchrun --nproc_per_node=4 --log-dir /tmp/mcore-single-weight-ut --tee 0:3 --redirects 3 -m pytest -s -v --confcutdir=tests/unit_tests/transformer/moe tests/unit_tests/transformer/moe/test_moe_single_grouped_weight_numerics.py

torchrun --nproc_per_node=4 --log-dir /tmp/mcore-grouped-dispatcher-ut --tee 0:3 --redirects 3 -m pytest -s -v --confcutdir=tests/unit_tests/transformer/moe tests/unit_tests/transformer/moe/test_grouped_tensor_dispatcher_numerics.py

E2E testing:

Model - Qwen3.5 VL SFT
Dispatcher - HybridEP
Experiements - bf16 cublas grouped gemm [single weight ON/OFF] vs. cuteDSL with TE op fuser mxfp8 [single weight ON/OFF] vs mxfp8 cublas grouped gemm [single weight ON/OFF]

image

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact @NVIDIA/mcore-oncall.

Issue tracking

For PRs from open-source community contributors:

  • New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
  • Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Linked issue:

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 @NVIDIA/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"

  1. When your PR is ready, click Ready for Review.
  2. 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.

@copy-pr-bot

copy-pr-bot Bot commented Jul 23, 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.

@zhongbozhu zhongbozhu mentioned this pull request Jul 29, 2026
Open
@zhongbozhu
zhongbozhu force-pushed the main_support_device_group_linear branch 2 times, most recently from 7294e71 to 1f4a49e Compare July 29, 2026 02:32
@zhongbozhu zhongbozhu mentioned this pull request Jul 31, 2026
Open
13 tasks
@zhongbozhu
zhongbozhu force-pushed the main_support_device_group_linear branch from 944cb4e to a78c12b Compare August 1, 2026 00:55
@zhongbozhu
zhongbozhu marked this pull request as ready for review August 1, 2026 09:12
@zhongbozhu
zhongbozhu requested review from a team as code owners August 1, 2026 09:12
@zhongbozhu

zhongbozhu commented Aug 1, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 1f647e6

zhongbozhu
zhongbozhu commented Aug 3, 2026
View reviewed changes
Comment thread megatron/core/transformer/moe/experts.py
)

@staticmethod
def _apply_packed_bias(intermediate_parallel, packed_bias, tokens_per_expert, permuted_probs):

@zhongbozhu zhongbozhu Aug 3, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Note: only models with routed expert bias will use this

@zhongbozhu

zhongbozhu commented Aug 4, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test d5c9aed

@zhongbozhu
zhongbozhu force-pushed the main_support_device_group_linear branch from d6d76b2 to f921404 Compare August 5, 2026 18:40
@zhongbozhu

zhongbozhu commented Aug 5, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test f921404

YangFei1990
YangFei1990 reviewed Aug 5, 2026
View reviewed changes

@YangFei1990 YangFei1990 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just for my understanding, mcore will pass moe_use_grouped_tensor / use_transformer_engine_op_fuser into TE, but it is TE's responsibility to pick the groupgemm backend, which is depending on hardware/dtype/envs, so it might not necessarily follow what exactly provided by mcore, is that right?

Comment thread megatron/core/transformer/moe/experts.py
# Some dispatchers have already padded each expert's token segment before the tokens
# reach this module:
# * router padding changes the routing map before dispatch;
# * HybridEP pads as part of its fused dispatch/permute operation;

@YangFei1990 YangFei1990 Aug 5, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: HybridEP/NCCL EP

Comment thread megatron/core/transformer/moe/experts.py
# The token buffer may already contain per-expert padding when padding was performed
# before expert compute:
# * router padding modified the routing map before dispatch;
# * HybridEP fused padding into dispatch/permute;

@YangFei1990 YangFei1990 Aug 5, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: same + NCCL EP

Comment thread tests/unit_tests/transformer/moe/test_grouped_tensor_dispatcher_numerics.py
assert torch.isfinite(hidden_states.grad).all()


class TestGroupedTensorDispatcherNumerics:

@YangFei1990 YangFei1990 Aug 5, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could we also add test cases for NCCL EP?

@zhongbozhu

zhongbozhu commented Aug 6, 2026

Copy link
Copy Markdown
Contributor Author

Just for my understanding, mcore will pass moe_use_grouped_tensor / use_transformer_engine_op_fuser into TE, but it is TE's responsibility to pick the groupgemm backend, which is depending on hardware/dtype/envs, so it might not necessarily follow what exactly provided by mcore, is that right?

If we have single weight, then fallback is banned. If we have discrete weight plus that certain recipes don't have fully implemented grouped kernels, it will fallback to CPU based m_split and launch multistream gemms.

@zhongbozhu
support device group linear
ceb69f1 
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>
@zhongbozhu
improve tests
b51e8f0 
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>
@zhongbozhu
fix unit test failure
9d74aee 
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>
@zhongbozhu
pass all tests
ff9437e 
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>
@zhongbozhu
format
b5d2bc2 
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>
@zhongbozhu
call te_general_gemm properly
4d88d61 
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>
@zhongbozhu
fix CI issues
c542386 
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>
@zhongbozhu
zhongbozhu force-pushed the main_support_device_group_linear branch from f921404 to c542386 Compare August 6, 2026 19:23
@zhongbozhu

zhongbozhu commented Aug 6, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test c542386

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

Reviewers

@YangFei1990 YangFei1990 YangFei1990 left review comments
@yaox12 yaox12 yaox12 approved these changes

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

complexity: medium

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@zhongbozhu @yaox12 @YangFei1990