Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Improve torch compile for fused moe #2327

Merged
merged 5 commits into from
Dec 3, 2024
Merged

Improve torch compile for fused moe #2327

merged 5 commits into from
Dec 3, 2024

Conversation

merrymercy
Copy link
Contributor

@merrymercy merrymercy commented Dec 3, 2024 •
edited
Loading

Following the discussion from #2278.

The final solution:

  • Only do torch.compile for moe layers at bs=1
  • Skip torch.compile for moe layers when bs > 1
  • Add dynamic=False in places where the shape should be known

@merrymercy merrymercy force-pushed the pr-improve-torch-moe branch from 59838b0 to 9bc7ef1 Compare December 3, 2024 05:47
@merrymercy merrymercy requested a review from hnyls2002 as a code owner December 3, 2024 06:13
@merrymercy merrymercy force-pushed the pr-improve-torch-moe branch from 5fa3529 to 88c361f Compare December 3, 2024 09:26
@merrymercy merrymercy changed the title improve torch compile for fused moe Improve torch compile for fused moe Dec 3, 2024
@merrymercy merrymercy mentioned this pull request Dec 3, 2024
Open
@merrymercy merrymercy merged commit 07ec07a into main Dec 3, 2024
3 of 16 checks passed
@merrymercy merrymercy deleted the pr-improve-torch-moe branch December 3, 2024 09:58
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@Ying1123 Ying1123 Awaiting requested review from Ying1123 Ying1123 is a code owner

@zhyncs zhyncs Awaiting requested review from zhyncs zhyncs is a code owner

@ispobock ispobock Awaiting requested review from ispobock ispobock is a code owner

@hnyls2002 hnyls2002 Awaiting requested review from hnyls2002 hnyls2002 is a code owner

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
1 participant
@merrymercy