Efficient Expert Weight Fusion for Moe deepseek v3#39150
Open
VassilyLombard wants to merge 2 commits into
Open
Efficient Expert Weight Fusion for Moe deepseek v3#39150VassilyLombard wants to merge 2 commits into
VassilyLombard wants to merge 2 commits into
Conversation
Author
|
Hi! First-time contributor here, happy to make any changes. I intended to do the second commit only, sorry for the confusion. Please let me know if you need anything else from me! |
ArthurZucker
reviewed
Jul 9, 2025
Comment on lines
+181
to
+183
| w1 = torch.stack([e.gate_proj.weight for e in self.experts], dim=0) | ||
| w2 = torch.stack([e.down_proj.weight for e in self.experts], dim=0) | ||
| w3 = torch.stack([e.up_proj.weight for e in self.experts], dim=0) |
Collaborator
There was a problem hiding this comment.
there is no way this is not super costly no? Also have you ran this with TP? we need to make sure both pathes work!
| expert_mask = torch.nn.functional.one_hot(topk_indices, num_classes=len(self.experts)) | ||
| expert_mask = expert_mask.permute(2, 0, 1) | ||
|
|
||
| for expert_idx in range(len(self.experts)): |
Collaborator
There was a problem hiding this comment.
we don't have to go through all the export tho! we should only go through the ones that are hit!
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
This PR answers a call for contribution. It introduces a fully vectorized, efficient Mixture-of-Experts (MoE) implementation for the DeepseekV3 model, specifically in the class DeepseekV3Moe. The new approach eliminates the loop over experts, improving inference and training speed, especially for large numbers of experts (256 here!).
Key Changes
Instead of iterating over experts, all expert weights are stacked into tensors, and expert selection is performed using indexing and batched matrix multiplications.
For each token, the top-k expert indices and routing weights are computed by the router. Inputs are repeated and routed to the selected experts in a single batched operation.
All expert computations (linear projections and activations) are performed in parallel for all tokens and their assigned experts.
The outputs from all routed experts are weighted and summed back to the original token positions, producing the final MoE output without explicit loop.
The implementation assumes all experts are instances of the same class and have the same architecture, enabling parameter stacking and efficient computation.
No additional dependencies are required. @ArthurZucker