Skip to content

fix(moe): keep precision-overridden grouped experts off the op-fuser path - #7212

Merged
fanshiqing merged 1 commit into
NVIDIA:mainfrom
fanshiqing:shiqingf/fix_gmm_fusion_with_precision_override
Sep 10, 2026
Merged

fanshiqing merged 1 commit into
NVIDIA:mainfrom
fanshiqing:shiqingf/fix_gmm_fusion_with_precision_override

Conversation

@fanshiqing

@fanshiqing fanshiqing commented Sep 10, 2026 •

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

What does this PR do?

fix(moe): keep precision-overridden grouped experts off the op-fuser path

  • Bug:
    • The op-fuser's grouped-MLP kernels are FP8/NVFP4-only and take their recipe
      from the global autocast state, so a --te-precision-config-file override never
      reached them.
    • Thus the layer quantized anyway: silently under plain TE, fatally under
      GTP, whose backward hands the kernel an unquantized weight (AttributeError on
      '_columnwise_data').
  • Fix: gate _with_fused_impl on both linears' will_execute_quantized(). Fusion spans
    fc1 and fc2, so either one opting out disables it.

Before the fix: casting+mxfp8-gmm-fusion

image

After the fix: fall back to unfused gmm in bf16

image

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

…path

- Bug: the op-fuser's grouped-MLP kernels are FP8/NVFP4-only and take their recipe
  from the global autocast state, so a --te-precision-config-file override never
  reached them. The layer quantized anyway: silently under plain TE, fatally under
  GTP, whose backward hands the kernel an unquantized weight (AttributeError on
  '_columnwise_data').
- Fix: gate _with_fused_impl on both linears' will_execute_quantized(). Fusion spans
  fc1 and fc2, so either one opting out disables it.

Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
@fanshiqing fanshiqing self-assigned this Sep 10, 2026
@fanshiqing
fanshiqing requested review from a team as code owners September 10, 2026 15:14
@copy-pr-bot

copy-pr-bot Bot commented Sep 10, 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.

@svcnvidia-nemo-ci
svcnvidia-nemo-ci marked this pull request as draft September 10, 2026 15:14
@github-actions

Copy link
Copy Markdown
Contributor

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:

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

See the contribution guide for more details.

@fanshiqing
fanshiqing marked this pull request as ready for review September 10, 2026 15:14
@fanshiqing

Copy link
Copy Markdown
Member Author

/ok to test 9f0564b

@fanshiqing

Copy link
Copy Markdown
Member Author

/claude review

@claude

This comment was marked as low quality.

@fanshiqing
fanshiqing requested a review from kvareddy September 10, 2026 16:01

@zhongbozhu zhongbozhu 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.

LGTM

@svcnvidia-nemo-ci svcnvidia-nemo-ci added the Final Review PR is in the "final review" stage label Sep 10, 2026
@svcnvidia-nemo-ci svcnvidia-nemo-ci removed the Final Review PR is in the "final review" stage label Sep 10, 2026
@svcnvidia-nemo-ci svcnvidia-nemo-ci added the Approved All necessary approvals have been made label Sep 10, 2026
@fanshiqing
fanshiqing added this pull request to the merge queue Sep 10, 2026
@nemo-automation-bot

Copy link
Copy Markdown

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/34522758319

Merged via the queue into NVIDIA:main with commit 8962453 Sep 10, 2026
120 of 122 checks passed
@fanshiqing
fanshiqing deleted the shiqingf/fix_gmm_fusion_with_precision_override branch September 10, 2026 21:09
@fanshiqing fanshiqing changed the title fix(moe): keep precision-overridden grouped experts off the op-fuser … fix(moe): keep precision-overridden grouped experts off the op-fuser path Sep 10, 2026
@fanshiqing fanshiqing linked an issue Sep 11, 2026 that may be closed by this pull request
@fanshiqing fanshiqing removed a link to an issue Sep 11, 2026
@fanshiqing fanshiqing linked an issue Sep 19, 2026 that may be closed by this pull request
xuwchen pushed a commit to xuwchen/Megatron-LM that referenced this pull request Oct 14, 2026
NVIDIA#7212)

Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
(cherry picked from commit 8962453)
Signed-off-by: xuwenc <xuwenc@nvidia.com>
xuwchen pushed a commit to xuwchen/Megatron-LM that referenced this pull request Oct 14, 2026
NVIDIA#7212)

Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
(cherry picked from commit 8962453)
Signed-off-by: xuwenc <xuwenc@nvidia.com>

This branch was previously deployed

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

Labels

Approved All necessary approvals have been made complexity: low nemotron

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[feat] GTP+MTP

7 participants