Skip to content

[Dev] JIT for MoE router and preprocess - #1918

Merged
yaox12 merged 8 commits into
NVIDIA:devfrom
yaox12:xiny/jit_for_moe
Oct 27, 2025
Merged

[Dev] JIT for MoE router and preprocess#1918
yaox12 merged 8 commits into
NVIDIA:devfrom
yaox12:xiny/jit_for_moe

Conversation

@yaox12

@yaox12 yaox12 commented Oct 24, 2025

Copy link
Copy Markdown
Member

PR to the main branch: #1919

  • Enable JIT (torch.compile) for expert bias and MoE preprocess.
  • Let force load balancing compatible with CUDA Graph.

Signed-off-by: Xin Yao <xiny@nvidia.com>
Signed-off-by: Xin Yao <xiny@nvidia.com>
Signed-off-by: Xin Yao <xiny@nvidia.com>
@yaox12
yaox12 requested review from a team as code owners October 24, 2025 07:32
@copy-pr-bot

copy-pr-bot Bot commented Oct 24, 2025

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.

@yaox12
yaox12 requested a review from hxbai October 24, 2025 07:33
@yaox12

yaox12 commented Oct 24, 2025

Copy link
Copy Markdown
Member Author

/ok to test cd17d4e

@hxbai

hxbai commented Oct 24, 2025

Copy link
Copy Markdown
Contributor

LGTM.

Comment thread megatron/core/transformer/moe/moe_utils.py Outdated
@yaox12

yaox12 commented Oct 24, 2025

Copy link
Copy Markdown
Member Author

/ok to test 41d4226

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

yaox12 commented Oct 27, 2025

Copy link
Copy Markdown
Member Author

/ok to test 12ff69b

@yaox12

yaox12 commented Oct 27, 2025

Copy link
Copy Markdown
Member Author

@Victarry @NVIDIA/core-nemo Can I get your approval?

@Victarry Victarry left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

LGTM!

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

Labels

dev branch Dev branch related issues and development module: moe

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants