Repository navigation
[Docs] Add Mixture of Experts guide #3494
There are no files selected for viewing
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
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
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
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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,30 @@ | ||
| # Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. | ||
| # | ||
| # See LICENSE for license information. | ||
|
|
||
| # START_GROUPED_LINEAR_JAX | ||
| import jax.numpy as jnp | ||
| from transformer_engine.jax import dense as te_dense | ||
|
|
||
| # x: [sum(group_sizes), hidden_size], expert-contiguous tokens | ||
| # kernel: [num_experts, hidden_size, ffn_hidden_size], stacked per-expert weights | ||
| # bias: [num_experts, ffn_hidden_size], stacked per-expert biases | ||
| # group_sizes: [num_experts] int array; group_sizes[i] is the number of routed | ||
| # tokens for expert i | ||
| split_indices = jnp.cumsum(group_sizes)[:-1] | ||
| x_by_expert = jnp.split(x, split_indices, axis=0) | ||
|
|
||
| # Baseline: one matmul per expert. | ||
| loop_out = jnp.concatenate( | ||
|
Collaborator
There was a problem hiding this comment. Thanks for adding JAX documentation as well @pggPL! 🙌 I've reviewed and left a few comments. Let me know what you think. Thanks! |
||
| [x_i @ kernel_i + bias_i for x_i, kernel_i, bias_i in zip(x_by_expert, kernel, bias)], | ||
| axis=0, | ||
| ) | ||
|
|
||
| # Transformer Engine: one grouped dense call. group_sizes is a device array. | ||
| grouped_out = te_dense.grouped_dense( | ||
|
Collaborator
There was a problem hiding this comment. @pggPL How do we want to handle group alignment in the docs? This does support MXFP8 but we require group sizes to be aligned to a multiple of 128 for MXFP8, so if a user passes in un-aligned groups we can get IMA, illegal instruction, or incorrect results For now, I've been guiding users to our more monolithic MoE block to avoid this complexity. But the same constraints apply to TE/PyTorch, so if you've found a better way to explain the nuances of this alignment, let me know and I'm open to adding it. The on-device group size alignment is difficult since we can't assert it on Host without introducing runtime overhead |
||
| x, | ||
| kernel, | ||
| group_sizes=group_sizes, | ||
| bias=bias, | ||
| ) | ||
| # END_GROUPED_LINEAR_JAX | ||
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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,43 @@ | ||
| # Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. | ||
| # | ||
| # See LICENSE for license information. | ||
|
|
||
| # START_GROUPED_LINEAR_PYTORCH | ||
| import torch | ||
| import transformer_engine.pytorch as te | ||
|
|
||
| # x: [sum(m_splits), hidden_size], expert-contiguous tokens | ||
| # m_splits: list[int] of length num_experts; m_splits[i] is the number | ||
| # of routed tokens for expert i | ||
| # torch_experts: list[torch.nn.Linear] of length num_experts, one per expert | ||
| # (used only by the baseline loop below) | ||
| x_by_expert = torch.split(x, m_splits, dim=0) | ||
|
|
||
| # Baseline: one Linear call per expert. | ||
| loop_out = torch.cat( | ||
| [expert(x_i) for expert, x_i in zip(torch_experts, x_by_expert)], | ||
| dim=0, | ||
| ) | ||
|
|
||
| # Transformer Engine: one grouped linear call. | ||
| grouped_linear = te.GroupedLinear( | ||
| num_experts, | ||
| hidden_size, | ||
| ffn_hidden_size, | ||
| bias=True, | ||
| params_dtype=torch.bfloat16, | ||
| ).cuda() | ||
| grouped_out = grouped_linear(x, m_splits) | ||
|
|
||
| # Keep token counts on the device with use_grouped_tensor=True. | ||
| grouped_linear = te.GroupedLinear( | ||
| num_experts, | ||
| hidden_size, | ||
| ffn_hidden_size, | ||
| bias=True, | ||
| params_dtype=torch.bfloat16, | ||
| use_grouped_tensor=True, | ||
| ).cuda() | ||
| m_splits_dev = torch.tensor(m_splits, dtype=torch.int64, device="cuda") | ||
| grouped_out = grouped_linear(x, m_splits_dev) | ||
| # END_GROUPED_LINEAR_PYTORCH |
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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,23 @@ | ||
| # Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. | ||
| # | ||
| # See LICENSE for license information. | ||
|
|
||
| # START_GROUPED_MLP_PYTORCH | ||
| import transformer_engine.pytorch as te | ||
|
|
||
| # FC1 produces gate and value features interleaved in blocks of 32. | ||
| expert_mlp = te.ops.Sequential( | ||
| te.ops.GroupedLinear(num_experts, hidden_size, 2 * ffn_hidden_size), | ||
| te.ops.ScaledSwiGLU(glu_interleave_size=32), | ||
| te.ops.GroupedLinear(num_experts, ffn_hidden_size, hidden_size, scale_bias=True), | ||
| ) | ||
|
pggPL marked this conversation as resolved.
|
||
|
|
||
| # Dispatch outputs: expert-contiguous tokens, aligned counts, and routing weights. | ||
| expert_out = expert_mlp( | ||
| permuted, | ||
| tokens_per_expert, | ||
| permuted_probs, | ||
| tokens_per_expert, | ||
| permuted_probs, | ||
| ) | ||
| # END_GROUPED_MLP_PYTORCH | ||
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Oops, something went wrong.
Oops, something went wrong.
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.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I don't think we should expose these API now.
I would rather deprecate them and expose the permutation ops alone instead, in case there are users interested in the permutation alone.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I'm not sure what you mean by "permutation ops". Can you please elaborate on that?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I mean permute alone, but not
token_dispatch/token_combine.It's confusing for users seeing
ep.ep_dispatchvspermutation.token_dispatch