Skip to content

Relax TE Op Fuser Env Var Assert NVTE_CUTEDSL_FUSED_GROUPED_MLP - #7499

Merged
zhongbozhu merged 2 commits into
NVIDIA:mainfrom
zhongbozhu:relax_te_op_fuser_checker
Oct 8, 2026
Merged

Relax TE Op Fuser Env Var Assert NVTE_CUTEDSL_FUSED_GROUPED_MLP#7499
zhongbozhu merged 2 commits into
NVIDIA:mainfrom
zhongbozhu:relax_te_op_fuser_checker

Conversation

@zhongbozhu

@zhongbozhu zhongbozhu commented Sep 19, 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?

TE should remove these env vars first: NVIDIA/TransformerEngine#3584

⚠️ 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
  • If this PR adds or changes a GPU kernel (Triton, jit_fuser/torch.compile, CUDA extension, TE or external-library dispatch, or a scatter/index accumulation), I have added or updated its bit-exact determinism test and registered it in tests/unit_tests/determinism/kernels/manifest.py (guide)
  • 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.

@zhongbozhu
relax te opfuser checker
17e943a 
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>
@zhongbozhu
zhongbozhu requested review from a team as code owners September 19, 2026 00:07
@copy-pr-bot

copy-pr-bot Bot commented Sep 19, 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 19, 2026 00:07
@github-actions

github-actions Bot commented Sep 19, 2026

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.

@zhongbozhu
zhongbozhu marked this pull request as ready for review October 1, 2026 22:40
@svcnvidia-nemo-ci svcnvidia-nemo-ci added the Final Review PR is in the "final review" stage label Oct 5, 2026
jiemingz
jiemingz reviewed Oct 6, 2026
View reviewed changes
Comment thread megatron/core/transformer/moe/experts.py
return False
# TE can execute GroupedLinear -> activation -> GroupedLinear without a fused kernel.
# Leave CuTeDSL enablement and layout checks to TE's fusion matcher so disabling
# NVTE_CUTEDSL_FUSED_GROUPED_MLP does not disable the basic op-fuser path.

@jiemingz jiemingz Oct 6, 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.

does this need to be TE version checked for back compat?

@zhongbozhu zhongbozhu Oct 6, 2026 •
edited
Loading

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.

TE always has the mechanism to fallback to unfused when certain fusion is not available. For legacy scripts with this env var turned on, they won't need any change, it's just that future TE won't read this env var anymore so it's NO-OP. For future scripts, they don't need to set this env var either because it will be automatically on.

@zhongbozhu
zhongbozhu enabled auto-merge October 6, 2026 19:27
@zhongbozhu

zhongbozhu commented Oct 6, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 17e943a

@zhongbozhu

zhongbozhu commented Oct 6, 2026

Copy link
Copy Markdown
Contributor Author

Adding determinism-exempt since this PR literally just removes an env var gate

@zhongbozhu
fix CI
28dde28 
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>
zhongbozhu
zhongbozhu commented Oct 7, 2026
View reviewed changes
Comment thread tests/unit_tests/transformer/moe/test_grouped_mlp.py
monkeypatch.setattr(experts_module, "te", fake_te)
monkeypatch.setattr(experts_module, "HAVE_TE", True)
monkeypatch.setattr(experts_module, "is_te_min_version", lambda _: True)
monkeypatch.delenv("NVTE_CUTEDSL_FUSED_GROUPED_MLP", raising=False)

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

this env is gonna be removed in both Mcore checks & TE, so this test doesn't make sense anymore

@zhongbozhu

zhongbozhu commented Oct 7, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 28dde28

@zhongbozhu
zhongbozhu added this pull request to the merge queue Oct 8, 2026
@nemo-automation-bot

nemo-automation-bot Bot commented Oct 8, 2026

Copy link
Copy Markdown

🔄 Merge queue validation started!

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

Merged via the queue into NVIDIA:main with commit 7e197c9 Oct 8, 2026
173 of 175 checks passed
@zhongbozhu
zhongbozhu deleted the relax_te_op_fuser_checker branch October 8, 2026 23:00

This branch was successfully deployed

2 active deployments
public — 28dde28d Deployed Oct 7, 2026 by copy-pr-bot[bot] via build-test-publish-wheels / publish-wheels (megatron-core, amd64) #11002
test — 28dde28d Deployed Oct 7, 2026 by copy-pr-bot[bot] via cicd-wait-in-queue #23389
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@jiemingz jiemingz jiemingz approved these changes
@yaox12 yaox12 yaox12 approved these changes
@cuichenx cuichenx cuichenx approved these changes
@yashaswikarnati yashaswikarnati yashaswikarnati approved these changes

Assignees

No one assigned

Labels

complexity: low determinism-exempt Final Review PR is in the "final review" stage

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

6 participants

@zhongbozhu @yaox12 @cuichenx @jiemingz @yashaswikarnati @svcnvidia-nemo-ci