[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,33 @@ | ||
| # 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( | ||
| [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. | ||
| # On Blackwell, BF16 and MXFP8 inputs without bias run as a single grouped GEMM | ||
|
Collaborator
There was a problem hiding this comment. Small update, Hopper BF16 is now supported since this PR merged yesterday: #3083 |
||
| # with the group sizes kept on the device; other cases launch one GEMM per | ||
| # expert and copy group_sizes to the host first. | ||
| 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,47 @@ | ||
| # 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. By default one GEMM per expert | ||
| # is launched; m_splits is read on the host. | ||
| 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) | ||
|
|
||
| # Single grouped GEMM with the token counts on the device (no host sync, | ||
| # CUDA-graph capturable): opt in with use_grouped_tensor=True and pass | ||
| # m_splits as a CUDA int64 tensor. Falls back to per-expert GEMMs when the | ||
| # recipe / GPU / cuBLAS version does not support it. | ||
| 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,17 @@ | ||
| # 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 for the GLU. | ||
| expert_mlp = te.ops.Sequential( | ||
| te.ops.GroupedLinear(num_experts, hidden_size, 2 * ffn_hidden_size), | ||
| te.ops.ScaledSwiGLU(), # or ScaledClampedQGeGLU; ScaledSReLU for the unary variant | ||
| te.ops.GroupedLinear(num_experts, ffn_hidden_size, hidden_size), | ||
| ) | ||
|
|
||
| # The fuser selects GroupedMLP_CuTeGEMMGLU for supported configurations. | ||
| # FC1 fuses the activation; FC2 runs as a separate grouped GEMM. | ||
| # 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.
Thanks for adding JAX documentation as well @pggPL! 🙌
I've reviewed and left a few comments. Let me know what you think. Thanks!