[MoE][PyTorch] Add mask-based MoE permutation #1373
+1,616
−40
Conversation
This file contains 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
hxbai
changed the title
yaox12
reviewed
Dec 13, 2024
Signed-off-by: Hongxiao Bai <[email protected]>
for more information, see https://pre-commit.ci Signed-off-by: Hongxiao Bai <[email protected]>
Signed-off-by: Hongxiao Bai <[email protected]>
Signed-off-by: Hongxiao Bai <[email protected]>
hxbai
force-pushed
the
permute_fusion
branch
from
6160104 to
ca94d72
Compare
timmoon10
reviewed
Jan 8, 2025
| ] | ||
|
|
||
|
|
||
| class _moe_permute(torch.autograd.Function): | ||
| """functional Permute""" | ||
| class _moe_permute_indice_map(torch.autograd.Function): |
There was a problem hiding this comment.
Suggested change
| class _moe_permute_indice_map(torch.autograd.Function): | |
| class _moe_permute_index_map(torch.autograd.Function): |
We should make sure to use "index" in user-facing APIs like moe_permute/moe_unpermute.
| import warnings | ||
| from typing import Tuple | ||
| import torch | ||
|
|
||
| import transformer_engine_torch as tex | ||
| from .constants import TE_DType | ||
| from .float8_tensor import Float8Tensor | ||
| import transformer_engine.pytorch.triton.permutation as triton_permuataion |
There was a problem hiding this comment.
Nit:
Suggested change
| import transformer_engine.pytorch.triton.permutation as triton_permuataion | |
| import transformer_engine.pytorch.triton.permutation as triton_permutation |
Comment on lines
+292
to
+295
| if ctx.fp8: | ||
| assert isinstance( | ||
| permuted_act_grad, Float8Tensor | ||
| ), "Grad of the output must be in Float8Tensor type for FP8 moe_permute." |
There was a problem hiding this comment.
Couldn't we decouple FP8 in the forward and backward?
Suggested change
| if ctx.fp8: | |
| assert isinstance( | |
| permuted_act_grad, Float8Tensor | |
| ), "Grad of the output must be in Float8Tensor type for FP8 moe_permute." | |
| fp8 = isinstance(permuted_act_grad, Float8Tensor) | |
| if fp8: |
If there are no obstacles, we could also do the same thing for _moe_unpermute_mask_map and _moe_chunk_sort.
| # Results Check | ||
| # | ||
| ################################################################################################################################### | ||
| tols = dtype_tols(te_dtype) |
There was a problem hiding this comment.
Shouldn't we expect bit-wise exact results?
Suggested change
| tols = dtype_tols(te_dtype) | |
| tols = { "atol": 0, "rtol": 0 } |
| # Results Check | ||
| # | ||
| ################################################################################################################################### | ||
| tols = dtype_tols(te_dtype) |
There was a problem hiding this comment.
We should expect bit-wise exact results.
Suggested change
| tols = dtype_tols(te_dtype) | |
| tols = { "atol": 0, "rtol": 0 } |
phu0ngng
reviewed
Jan 10, 2025
| mask=(offset < num_tokens), | ||
| other=0, | ||
| ).to(tl.int64) | ||
| expert_token_cumsum = tl.cumsum(expert_token_mask) * expert_token_mask |
There was a problem hiding this comment.
An interesting way to exclude the zero token_mask. Happy to learn!
Comment on lines
+61
to
+67
| chunk_cumsum = tl.load( | ||
| row_id_map_ptr + pid_m * num_tokens + offset, mask=(offset < num_tokens), other=0 | ||
| ) | ||
|
|
||
| workspace_off = tl.arange(0, WORKSPACE_LOAD_WIDTH) | ||
| chunk_sums = tl.load(workspace_ptr + workspace_off, mask=workspace_off < chunk_idx) | ||
| chunk_cumsum = tl.where(chunk_cumsum == 0, -1, chunk_cumsum + tl.sum(chunk_sums) - 1) |
There was a problem hiding this comment.
These three names chuck_cumsum, chuck_sums, and chunk_cumsum are quite confusing.
If I understand it correctly, I suggest to rename them to:
chuck_cumsum->row_id_within_token_blockchuck_sums-> n_tokens_per_expertchuck_cumsum->row_id
In addition, I think we should move the -1 to the pass1 as it is the correction for the calculation of expert_token_cumsum, as:
expert_token_cumsum = (tl.cumsum(expert_token_mask) - 1) * expert_token_mask
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
Description
Add mask-based token permutation and local chunk permutation fused kernels. These kernels are implemented with OpenAI Triton.
Related commit in Megatron-LM NVIDIA/Megatron-LM@ac0474d
Fixes # (issue)
Type of change
Changes
Please list the changes introduced in this PR:
te.pytorch.permutation.moe_permuteandte.pytorch.permutation.moe_unpermutete.pytorch.permutation.moe_sort_chunks_by_indicesChecklist: