Fuse permute+pad and unpermute+unpad ops for FP8/FP4 precision

This can remove explicit padding/unpadding around GroupedMLP,
which improves throughput and reduces peak memory usage

Signed-off-by: xiaoxi-wangfj <690912414@qq.com>

Author: xiaoxi-wangfj

megatron/core/extensions/transformer_engine.py

@@ -2328,6 +2328,14 @@ def fused_apply_rotary_pos_emb_thd(
     fused_sort_chunks_by_index_with_probs = None
     fused_unpermute = None
 
+try:
+    from transformer_engine.pytorch.permutation import moe_permute_and_pad_with_probs
+
+    fused_permute_and_pad_with_probs = moe_permute_and_pad_with_probs
+
+except ImportError:
+    fused_permute_and_pad_with_probs = None
+
 try:
     from transformer_engine.pytorch.cross_entropy import parallel_cross_entropy
 

megatron/core/transformer/moe/moe_utils.py

@@ -19,6 +19,7 @@
         fused_compute_score_for_moe_aux_loss,
         fused_moe_aux_loss,
         fused_permute,
+        fused_permute_and_pad_with_probs,
         fused_permute_with_probs,
         fused_sort_chunks_by_index,
         fused_sort_chunks_by_index_with_probs,
@@ -227,6 +228,8 @@ def permute(
     num_out_tokens: Optional[int] = None,
     fused: bool = False,
     drop_and_pad: bool = False,
+    tokens_per_expert: Optional[torch.Tensor] = None,
+    align_size: int = -1,
 ):
     """Permute the tokens and probs based on the mask.
     Tokens with the same designated expert will be grouped together.
@@ -257,11 +260,18 @@ def permute(
         return permuted_input, None, sorted_indices
 
     if fused and probs is not None:
-        if not HAVE_TE or fused_permute_with_probs is None:
-            raise ValueError(
-                "fused_permute_with_probs is not available. Please install TE >= 2.1.0."
+        if tokens_per_expert is not None and align_size > 0:
+            return fused_permute_and_pad_with_probs(
+                tokens, probs, routing_map, tokens_per_expert, align_size
+            )
+        else:
+            if not HAVE_TE or fused_permute_with_probs is None:
+                raise ValueError(
+                    "fused_permute_with_probs is not available. Please install TE >= 2.1.0."
+                )
+            return fused_permute_with_probs(
+                tokens, probs, routing_map, num_out_tokens=num_out_tokens
             )
-        return fused_permute_with_probs(tokens, probs, routing_map, num_out_tokens=num_out_tokens)
 
     num_tokens, hidden = tokens.shape
     num_experts = routing_map.shape[1]
@@ -315,6 +325,7 @@ def unpermute(
     routing_map: torch.Tensor = None,
     fused: bool = False,
     drop_and_pad: bool = False,
+    pad_offsets: Optional[torch.Tensor] = None,
 ):
     """
     Restore the original order of tokens after permutation. If probs are provided, it
@@ -344,7 +355,11 @@ def unpermute(
         if not HAVE_TE or fused_unpermute is None:
             raise ValueError("fused_unpermute is not available. Please install TE >= 2.1.0.")
         return fused_unpermute(
-            permuted_tokens, sorted_indices, merging_probs=probs, restore_shape=restore_shape
+            permuted_tokens,
+            sorted_indices,
+            merging_probs=probs,
+            restore_shape=restore_shape,
+            pad_offsets=pad_offsets,
         )
 
     _, hidden = restore_shape

megatron/core/transformer/moe/token_dispatcher.py

@@ -1258,13 +1258,29 @@ def get_permuted_hidden_states_by_experts(self, hidden_states: torch.Tensor) ->
 
         self.hidden_shape_before_permute = hidden_states.shape
         assert self.dispatched_probs.dtype == torch.float32, "DeepEP only supports float32 probs"
-        hidden_states, permuted_probs, self.reversed_mapping_for_combine = permute(
-            hidden_states,
-            self.dispatched_routing_map,
-            probs=self.dispatched_probs,
-            num_out_tokens=self.tokens_per_expert.sum().item(),
-            fused=self.permute_fusion,
-        )
+        if self.config.moe_permute_padding_for_quantization:
+            (
+                hidden_states,
+                permuted_probs,
+                self.reversed_mapping_for_combine,
+                self.pad_offsets,
+                self.tokens_per_expert,
+            ) = permute(
+                hidden_states,
+                self.dispatched_routing_map,
+                probs=self.dispatched_probs,
+                fused=self.permute_fusion,
+                tokens_per_expert=self.tokens_per_expert,
+                align_size=get_align_size_for_quantization(self.config),
+            )
+        else:
+            hidden_states, permuted_probs, self.reversed_mapping_for_combine = permute(
+                hidden_states,
+                self.dispatched_routing_map,
+                probs=self.dispatched_probs,
+                num_out_tokens=self.tokens_per_expert.sum().item(),
+                fused=self.permute_fusion,
+            )
         if self.router_dtype == "fp64":
             permuted_probs = permuted_probs.to(torch.float64)
         return hidden_states, permuted_probs
@@ -1276,6 +1292,9 @@ def get_restored_hidden_states_by_experts(self, hidden_states: torch.Tensor) ->
             restore_shape=self.hidden_shape_before_permute,
             routing_map=self.dispatched_routing_map,
             fused=self.permute_fusion,
+            pad_offsets=(
+                self.pad_offsets if self.config.moe_permute_padding_for_quantization else None
+            ),
         )
         return hidden_states
 

megatron/core/transformer/transformer_config.py

@@ -496,6 +496,12 @@ class TransformerConfig(ModelParallelConfig):
     """[Compatibility alias for moe_router_padding_for_quantization]
     Enabling this will also enable moe_router_padding_for_quantization."""
 
+    moe_permute_padding_for_quantization: Optional[bool] = False
+    """Enable padding during MoE token permutation and corresponding unpadding during unpermutation
+    so that the number of tokens in each expert's permuted block is aligned to a multiple of 16 / 32
+    for quantized precisions such as FP8 and FP4. This can remove explicit padding/
+    unpadding around GroupedMLP kernels, which improves throughput and reduces peak memory usage."""
+
     moe_router_num_groups: Optional[int] = None
     """Number of groups to divide experts into for group-limited routing.
     When using group-limited routing:
@@ -1419,6 +1425,22 @@ def __post_init__(self):
                     "moe_router_padding_for_quantization."
                 )
 
+        if self.moe_permute_padding_for_quantization:
+            if self.fp8 is None and self.fp4 is None:
+                raise ValueError(
+                    "moe_permute_padding_for_quantization requires a quantized precision recipe(e.g., fp8 or fp4) to be enabled."
+                )
+
+            if not self.moe_permute_fusion:
+                raise ValueError(
+                    "moe_permute_padding_for_quantization currently requires fused permute."
+                )
+
+            from megatron.core.transformer.moe.moe_utils import fused_permute_and_pad_with_probs
+
+            if fused_permute_and_pad_with_probs is None:
+                raise ValueError("fused_permute_and_pad_with_probs is not available. Please install TE >= 2.12.0.")
+
         if (
             self.moe_router_topk == 1
             and self.moe_router_score_function == "softmax"

megatron/training/arguments.py

@@ -3207,6 +3207,10 @@ def _add_moe_args(parser):
     group.add_argument('--moe-router-padding-for-fp8', action='store_true',
                        help='[Compatibility alias for --moe-router-padding-for-quantization] '
                        'Enabling this will also enable --moe-router-padding-for-quantization.')
+    group.add_argument('--moe-permute-padding-for-quantization', action='store_true',
+                       help='Enable padding during MoE token permutation (and unpadding during unpermutation) '
+                       'so that the number of tokens in each expert permuted block is aligned to a multiple of 16/32 for FP8/FP4 precision. '
+                       'This can remove explicit padding/unpadding around GroupedMLP, which improves throughput and reduces peak memory usage.')
     group.add_argument('--moe-aux-loss-coeff', type=float, nargs='+', default=0.0,
                        help='Scaling coefficient for the aux loss: a starting value of 1e-2 is recommended.')
     group.add_argument('--moe-z-loss-coeff', type=float, default=None,