From c49ffc102b5224e41ad1e48a44fa79868e17d95e Mon Sep 17 00:00:00 2001
From: HaochenYuan <haocheny@nvidia.com>
Date: Thu, 6 Nov 2025 02:35:23 -0800
Subject: [PATCH 01/32] GPG sign off

---
 megatron/core/models/gpt/gpt_model.py         |   2 +
 megatron/core/transformer/moe/moe_layer.py    |  26 +++-
 megatron/core/transformer/moe/moe_utils.py    | 108 ++++++++++++-
 megatron/core/transformer/moe/router.py       | 141 +++++++++++++----
 .../core/transformer/transformer_block.py     |   4 +-
 .../core/transformer/transformer_layer.py     |  67 +++++++--
 .../transformer/moe/test_aux_loss.py          | 142 ++++++++++++++++++
 .../transformer/moe/test_routers.py           |  36 +++++
 8 files changed, 469 insertions(+), 57 deletions(-)

diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py
index cd13a598d9e..3f9f34f9b59 100644
--- a/megatron/core/models/gpt/gpt_model.py
+++ b/megatron/core/models/gpt/gpt_model.py
@@ -445,6 +445,7 @@ def forward(
         *,
         inference_params: Optional[BaseInferenceContext] = None,
         loss_mask: Optional[Tensor] = None,
+        padding_mask: Optional[Tensor] = None,
     ) -> Tensor:
         """Forward function of the GPT Model This function passes the input tensors
         through the embedding layer, and then the decoder and finally into the post
@@ -486,6 +487,7 @@ def forward(
             rotary_pos_cos_sin=rotary_pos_cos_sin,
             packed_seq_params=packed_seq_params,
             sequence_len_offset=sequence_len_offset,
+            padding_mask=padding_mask,
             **(extra_block_kwargs or {}),
         )
 
diff --git a/megatron/core/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py
index 893b2e7b99a..5c8ca3181ab 100644
--- a/megatron/core/transformer/moe/moe_layer.py
+++ b/megatron/core/transformer/moe/moe_layer.py
@@ -178,13 +178,19 @@ def __init__(
         self.cudagraph_tensor_store = MoECudaGraphTensorStore()
 
     @maybe_skip_or_early_return_by_cudagraph("route")
-    def route(self, hidden_states: torch.Tensor):
+    def route(self, hidden_states: torch.Tensor, padding_mask: Optional[torch.Tensor] = None):
         """Compute token routing for preprocessing.
 
         This method uses the router to determine which experts to send each token to,
         producing routing probabilities and a mapping.
+        
+        Args:
+            hidden_states (torch.Tensor): Input tensor.
+            padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
+                                                   Shape in [seq_length, bsz]. True for valid tokens,
+                                                   False for padding tokens. Defaults to None.
         """
-        probs, routing_map = self.router(hidden_states)
+        probs, routing_map = self.router(hidden_states, padding_mask=padding_mask)
         return probs, routing_map
 
     @maybe_skip_or_early_return_by_cudagraph("preprocess")
@@ -270,7 +276,7 @@ def combine(self, output: torch.Tensor, shared_expert_output: Optional[torch.Ten
             output = output + shared_expert_output
         return output
 
-    def forward(self, hidden_states: torch.Tensor):
+    def forward(self, hidden_states: torch.Tensor, padding_mask: Optional[torch.Tensor] = None):
         """Forward pass for the MoE layer.
 
         The forward pass comprises four main steps:
@@ -280,7 +286,10 @@ def forward(self, hidden_states: torch.Tensor):
         4. Combine: The outputs from the experts are combined and returned.
 
         Args:
-            hidden_states (torch.Tensor): The input tensor to the MoE layer.
+            hidden_states (torch.Tensor): The input tensor to the MoE layer. 
+            padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
+                                                   Shape in [seq_length, bsz]. True for valid tokens,
+                                                   False for padding tokens. Defaults to None.
 
         Returns:
             A tuple containing the output tensor and the MLP bias, if any.
@@ -292,10 +301,10 @@ def forward(self, hidden_states: torch.Tensor):
             )
 
         # MoE forward: route -> dispatch -> compute -> combine
-        def custom_forward(hidden_states):
+        def custom_forward(hidden_states, padding_mask=None):
             try:
                 shared_expert_output = self.shared_experts_compute(hidden_states)
-                probs, routing_map = self.route(hidden_states)
+                probs, routing_map = self.route(hidden_states, padding_mask=padding_mask)
                 hidden_states, probs, residual = self.preprocess(hidden_states, probs, routing_map)
             except MoECudaGraphPartialCaptureSignal as e:
                 # This signal is raised from the maybe_skip_or_early_return_by_cudagraph decorator.
@@ -318,11 +327,12 @@ def custom_forward(hidden_states):
                     tensor_parallel.random.get_cuda_rng_tracker,
                     parallel_state.get_tensor_model_parallel_group(),
                     hidden_states,
+                    padding_mask,
                 )
             else:
-                outputs = tensor_parallel.checkpoint(custom_forward, False, hidden_states)
+                outputs = tensor_parallel.checkpoint(custom_forward, False, hidden_states, padding_mask)
         else:
-            outputs = custom_forward(hidden_states)
+            outputs = custom_forward(hidden_states, padding_mask)
 
         return outputs
 
diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py
index 5a0793ef5b9..5f5c59f5b2a 100644
--- a/megatron/core/transformer/moe/moe_utils.py
+++ b/megatron/core/transformer/moe/moe_utils.py
@@ -34,6 +34,69 @@
 _MOE_LAYER_WISE_LOGGING_TRACKER = {}
 
 
+def compute_tokens_per_expert_with_padding(
+    routing_map: torch.Tensor,
+    padding_mask: Optional[torch.Tensor] = None,
+    reshape_for_seq_aux: bool = False,
+    seq_length: Optional[int] = None,
+    bsz: Optional[int] = None,
+    num_experts: Optional[int] = None,
+):
+    """Compute tokens_per_expert and total_num_tokens with optional padding mask.
+    This function provides a unified way to compute token counts across different aux loss types.
+    
+    Args:
+        routing_map (torch.Tensor): Token to expert routing map.
+            - For aux_loss/global_aux_loss: shape [num_tokens, num_experts]
+            - For seq_aux_loss: shape [num_tokens, num_experts] but will be reshaped
+        padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
+            Shape [num_tokens]. True for valid tokens, False for padding tokens.
+        reshape_for_seq_aux (bool): If True, reshape routing_map for seq_aux_loss computation.
+        seq_length (int, optional): Sequence length, required when reshape_for_seq_aux=True.
+        bsz (int, optional): Batch size, required when reshape_for_seq_aux=True.
+        num_experts (int, optional): Number of experts, required when reshape_for_seq_aux=True.
+    
+    Returns:
+        tuple: (tokens_per_expert, num_valid_tokens)
+            - tokens_per_expert (torch.Tensor): Number of tokens per expert, shape [num_experts] 
+                or [bsz * num_experts] for seq_aux_loss
+            - num_valid_tokens (torch.Tensor or int): Number of valid (non-padding) tokens
+    """
+    if reshape_for_seq_aux:
+        # seq aux loss
+        assert seq_length is not None and bsz is not None and num_experts is not None, \
+            "seq_length, bsz, and num_experts must be provided when reshape_for_seq_aux=True"
+        
+        if padding_mask is not None:
+            # Reshape padding_mask to [seq_length, bsz]
+            padding_mask_reshaped = padding_mask.reshape(seq_length, bsz)
+            # Expand to match routing_map after reshape [seq_length, bsz * num_experts]
+            mask_expanded = padding_mask_reshaped.unsqueeze(-1).expand(
+                -1, -1, num_experts
+            ).reshape(seq_length, -1)
+            routing_map_masked = routing_map.reshape(seq_length, -1) & mask_expanded
+            tokens_per_expert = routing_map_masked.sum(dim=0)
+            # Count valid tokens only
+            num_valid_tokens = padding_mask.sum()
+        else:
+            tokens_per_expert = routing_map.reshape(seq_length, -1).sum(dim=0)
+            num_valid_tokens = routing_map.shape[0]
+    else:
+        # aux_loss or global_aux_loss
+        if padding_mask is not None:
+            # routing_map: [num_tokens, num_experts], padding_mask: [num_tokens]
+            mask_expanded = padding_mask.unsqueeze(-1)
+            routing_map_masked = routing_map & mask_expanded
+            tokens_per_expert = routing_map_masked.sum(dim=0)
+            # Count valid tokens only
+            num_valid_tokens = padding_mask.sum()
+        else:
+            tokens_per_expert = routing_map.sum(dim=0)
+            num_valid_tokens = routing_map.shape[0]
+    
+    return tokens_per_expert, num_valid_tokens
+
+
 def switch_load_balancing_loss_func(
     probs: torch.Tensor,
     tokens_per_expert: torch.Tensor,
@@ -42,6 +105,7 @@ def switch_load_balancing_loss_func(
     num_experts: int,
     moe_aux_loss_coeff: float,
     fused: bool = False,
+    padding_mask: Optional[torch.Tensor] = None,
 ):
     """Calculate the auxiliary loss for load balancing.
     Refer to the Switch Transformer (https://arxiv.org/abs/2101.03961)
@@ -92,9 +156,18 @@ def switch_load_balancing_loss_func(
         topk (int): The number of experts selected for each token.
         num_experts (int): The number of experts.
         moe_aux_loss_coeff (float): The coefficient for the auxiliary loss.
+        padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
+                                               Shape in [num_tokens]. True for valid tokens,
+                                               False for padding tokens. Defaults to None.
     Returns:
         torch.Tensor: The auxiliary loss for load balancing.
     """
+    # Apply padding mask to probs if provided
+    if padding_mask is not None:
+        # padding_mask: [num_tokens], probs: [num_tokens, num_experts]
+        mask_expanded = padding_mask.unsqueeze(-1)
+        probs = probs * mask_expanded
+    
     if fused:
         if not HAVE_TE or fused_moe_aux_loss is None:
             raise ValueError("fused_moe_aux_loss is not available. Please install TE >= 2.7.0.")
@@ -114,18 +187,32 @@ def switch_load_balancing_loss_func(
     return aux_loss
 
 
-def z_loss_func(logits, z_loss_coeff):
+def z_loss_func(logits, z_loss_coeff, padding_mask: Optional[torch.Tensor] = None):
     """Encourages the router's logits to remain small to enhance stability.
     Please refer to the ST-MoE paper (https://arxiv.org/pdf/2202.08906.pdf) for details.
 
     Args:
         logits (torch.Tensor): The logits of the router.
+        z_loss_coeff (float): The coefficient for the z-loss.
+        padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
+                                               Shape in [num_tokens]. True for valid tokens,
+                                               False for padding tokens. Defaults to None.
 
     Returns:
         torch.Tensor: The logits after applying the z-loss.
     """
+    logsum = torch.logsumexp(logits, dim=-1)
+    z_loss_values = torch.square(logsum)
+    
+    if padding_mask is not None:
+        # Only compute z_loss for non-padding tokens
+        z_loss_values = z_loss_values * padding_mask
+        # Compute mean over valid tokens only
+        num_valid_tokens = padding_mask.sum()
+        z_loss = z_loss_values.sum() / torch.clamp(num_valid_tokens, min=1.0) * z_loss_coeff
+    else:
+        z_loss = torch.mean(z_loss_values) * z_loss_coeff
 
-    z_loss = torch.mean(torch.square(torch.logsumexp(logits, dim=-1))) * z_loss_coeff
     return z_loss
 
 
@@ -623,21 +710,27 @@ def compute_topk(scores, topk, num_groups=None, group_topk=None):
 
 
 def compute_routing_scores_for_aux_loss(
-    logits: torch.Tensor, topk: int, score_function: str, fused: bool = False
+    logits: torch.Tensor, topk: int, score_function: str, fused: bool = False, 
+    padding_mask: Optional[torch.Tensor] = None
 ):
     """Compute routing scores based on the score function.
 
     Args:
         logits (torch.Tensor): The logits tensor after gating, shape: [num_tokens, num_experts].
-
+        padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
+                                               Shape in [num_tokens]. True for valid tokens,
+                                               False for padding tokens. Defaults to None.
     Returns:
-        torch.Tensor: The normalized routing scores.
+        Tuple[torch.Tensor, torch.Tensor]: routing_map and scores.
     """
     if fused:
         if not HAVE_TE or fused_compute_score_for_moe_aux_loss is None:
             raise ValueError(
                 "fused_compute_score_for_moe_aux_loss is not available. Please install TE >= 2.6.0."
             )
+        # Note: fused implementation does not support padding_mask yet
+        if padding_mask is not None:
+            raise ValueError("Fused compute_routing_scores does not support padding_mask. Set fused=False.")
         return fused_compute_score_for_moe_aux_loss(
             logits=logits, topk=topk, score_function=score_function
         )
@@ -652,6 +745,11 @@ def compute_routing_scores_for_aux_loss(
 
     _, top_indices = torch.topk(scores, k=topk, dim=1)
     routing_map = torch.zeros_like(logits).int().scatter(1, top_indices, 1).bool()
+    
+    # Apply padding mask to scores if provided
+    if padding_mask is not None:
+        mask_expanded = padding_mask.unsqueeze(-1).to(scores.dtype)
+        scores = scores * mask_expanded
     return routing_map, scores
 
 
diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py
index 16fc9d9af8f..83d0f8c131e 100644
--- a/megatron/core/transformer/moe/router.py
+++ b/megatron/core/transformer/moe/router.py
@@ -1,7 +1,7 @@
 # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 from abc import ABC, abstractmethod
-from typing import Optional
+from typing import Optional, Union
 
 import torch
 
@@ -14,6 +14,7 @@
     apply_random_logits,
     apply_router_token_dropping,
     compute_routing_scores_for_aux_loss,
+    compute_tokens_per_expert_with_padding,
     router_gating_linear,
     save_to_aux_losses_tracker,
     sinkhorn,
@@ -268,17 +269,24 @@ def is_aux_loss_enabled(self) -> bool:
         return False
 
     def _apply_aux_loss(
-        self, probs: torch.Tensor, scores_for_aux_loss: torch.Tensor, routing_map: torch.Tensor
+        self, probs: torch.Tensor, scores_for_aux_loss: torch.Tensor, routing_map: torch.Tensor,
+        padding_mask: Optional[torch.Tensor] = None
     ):
         """Apply the auxiliary loss for the given scores and routing map."""
         aux_loss_coeff = self.get_aux_loss_coeff("aux_loss")
         if aux_loss_coeff == 0:
             return probs
-        tokens_per_expert = routing_map.sum(dim=0)
+        
+        # Use unified function to compute tokens_per_expert and num_tokens
+        tokens_per_expert, num_tokens = compute_tokens_per_expert_with_padding(
+            routing_map=routing_map,
+            padding_mask=padding_mask,
+            reshape_for_seq_aux=False,
+        )
+        
         tokens_per_expert = reduce_from_tensor_model_parallel_region(
             tokens_per_expert, self.tp_cp_group
         )
-        num_tokens = routing_map.shape[0]
         total_num_tokens = num_tokens * self.tp_cp_group.size()
 
         aux_loss = switch_load_balancing_loss_func(
@@ -289,9 +297,11 @@ def _apply_aux_loss(
             num_experts=self.config.num_moe_experts,
             moe_aux_loss_coeff=aux_loss_coeff,
             fused=self.config.moe_router_fusion,
+            padding_mask=padding_mask,
         )
         probs = self.attach_and_log_load_balancing_loss(
-            probs, aux_loss_coeff, aux_loss, "load_balancing_loss", self.tp_cp_group
+            probs, aux_loss_coeff, aux_loss, "load_balancing_loss", self.tp_cp_group,
+            valid_token_count=num_tokens
         )
         return probs
 
@@ -302,6 +312,7 @@ def _apply_seq_aux_loss(
         routing_map: torch.Tensor,
         seq_length: int,
         bsz: int,
+        padding_mask: Optional[torch.Tensor] = None,
     ):
         """Apply the sequence-level auxiliary loss for the given scores and routing map.
 
@@ -315,13 +326,23 @@ def _apply_seq_aux_loss(
             return probs
 
         scores_for_aux_loss = scores_for_aux_loss.reshape(seq_length, -1)
-        tokens_per_expert = routing_map.reshape(seq_length, -1).sum(dim=0)
+        
+        # Use unified function to compute tokens_per_expert and num_tokens
+        tokens_per_expert, num_tokens = compute_tokens_per_expert_with_padding(
+            routing_map=routing_map,
+            padding_mask=padding_mask,
+            reshape_for_seq_aux=True,
+            seq_length=seq_length,
+            bsz=bsz,
+            num_experts=self.config.num_moe_experts,
+        )
+        
+        total_num_tokens = num_tokens * self.tp_cp_group.size()
+        
         tokens_per_expert = reduce_from_tensor_model_parallel_region(
             tokens_per_expert, self.tp_cp_group
         )
 
-        total_num_tokens = seq_length * self.tp_cp_group.size()
-
         aux_loss = (
             switch_load_balancing_loss_func(
                 probs=scores_for_aux_loss,
@@ -331,23 +352,38 @@ def _apply_seq_aux_loss(
                 num_experts=self.config.num_moe_experts,
                 moe_aux_loss_coeff=seq_aux_loss_coeff,
                 fused=self.config.moe_router_fusion,
+                padding_mask=None,
             )
             / bsz
         )
+        # Calculate valid token count: seq_length for each batch element
+        if padding_mask is not None:
+            valid_token_count = padding_mask.sum()
+        else:
+            valid_token_count = seq_length * bsz
+        
         probs = self.attach_and_log_load_balancing_loss(
-            probs, seq_aux_loss_coeff, aux_loss, "seq_load_balancing_loss", self.tp_cp_group
+            probs, seq_aux_loss_coeff, aux_loss, "seq_load_balancing_loss", self.tp_cp_group,
+            valid_token_count=valid_token_count
         )
         return probs
 
     def _apply_global_aux_loss(
-        self, probs: torch.Tensor, scores_for_aux_loss: torch.Tensor, routing_map: torch.Tensor
+        self, probs: torch.Tensor, scores_for_aux_loss: torch.Tensor, routing_map: torch.Tensor,
+        padding_mask: Optional[torch.Tensor] = None
     ):
         """Apply the global auxiliary loss for the given scores and routing map."""
         global_aux_loss_coeff = self.get_aux_loss_coeff("global_aux_loss")
         if global_aux_loss_coeff == 0:
             return probs
 
-        tokens_per_expert = routing_map.sum(dim=0)
+        # Use unified function to compute tokens_per_expert and num_tokens
+        tokens_per_expert, num_tokens = compute_tokens_per_expert_with_padding(
+            routing_map=routing_map,
+            padding_mask=padding_mask,
+            reshape_for_seq_aux=False,
+        )
+        
         tokens_per_expert = reduce_from_tensor_model_parallel_region(
             tokens_per_expert, self.tp_dp_cp_group
         )
@@ -356,7 +392,6 @@ def _apply_global_aux_loss(
         self.ga_steps += 1
         averated_tokens_per_expert = self.global_tokens_per_expert / self.ga_steps
 
-        num_tokens = scores_for_aux_loss.shape[0]
         total_num_tokens = num_tokens * self.tp_dp_cp_group.size()
 
         global_aux_loss = switch_load_balancing_loss_func(
@@ -367,6 +402,7 @@ def _apply_global_aux_loss(
             num_experts=self.config.num_moe_experts,
             moe_aux_loss_coeff=global_aux_loss_coeff,
             fused=self.config.moe_router_fusion,
+            padding_mask=None,
         )
         probs = self.attach_and_log_load_balancing_loss(
             probs,
@@ -374,6 +410,7 @@ def _apply_global_aux_loss(
             global_aux_loss,
             "global_load_balancing_loss",
             self.tp_dp_cp_group,
+            valid_token_count=num_tokens,
         )
         return probs
 
@@ -384,8 +421,20 @@ def attach_and_log_load_balancing_loss(
         aux_loss: torch.Tensor,
         aux_loss_name: str,
         reduce_group: torch.distributed.ProcessGroup,
+        valid_token_count: Optional[Union[int, torch.Tensor]] = None,
     ):
-        """Attach aux loss function to activation and add to logging."""
+        """Attach aux loss function to activation and add to logging.
+        
+        Args:
+            activation (torch.Tensor): Activation tensor to attach the aux loss to.
+            aux_loss_coeff (float): Coefficient for the aux loss.
+            aux_loss (torch.Tensor): Computed aux loss.
+            aux_loss_name (str): Name of the aux loss for logging.
+            reduce_group (torch.distributed.ProcessGroup): Process group for reduction.
+            valid_token_count (int or torch.Tensor, optional): Number of valid tokens excluding 
+                padding tokens. Can be a Python int or a torch.Tensor (typically 0-d tensor).
+                If None, uses activation.shape[0]. Defaults to None.
+        """
         # TODO (zijiey): fix the per_layer_logging for MTP, currently it will incorrectly
         # add the aux loss logging value to other layer's since it is difficult to get the
         # correct layer_number for MTP. It does not affect the correctness of the calculation
@@ -408,17 +457,22 @@ def attach_and_log_load_balancing_loss(
             # which scales both the main_loss gradient and aux_loss gradient by
             # 1/(num_local_tokens * dp_size * num_micro_batches) in finalize_model_grads function.
             # To correct this scaling, we need to scale the aux_loss by num_local_tokens here.
-            activation = MoEAuxLossAutoScaler.apply(activation, aux_loss * activation.shape[0])
+            # Use valid_token_count (excluding padding) if provided, otherwise use total tokens.
+            num_tokens = valid_token_count if valid_token_count is not None else activation.shape[0]
+            activation = MoEAuxLossAutoScaler.apply(activation, aux_loss * num_tokens)
         else:
             activation = MoEAuxLossAutoScaler.apply(activation, aux_loss)
         return activation
 
-    def apply_z_loss(self, logits):
+    def apply_z_loss(self, logits, padding_mask: Optional[torch.Tensor] = None):
         """Encourages the router's logits to remain small to enhance stability.
         Please refer to the ST-MoE paper (https://arxiv.org/pdf/2202.08906.pdf) for details.
 
         Args:
             logits (torch.Tensor): The logits of the router.
+            padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
+                                                   Shape in [num_tokens]. True for valid tokens,
+                                                   False for padding tokens. Defaults to None.
 
         Returns:
             torch.Tensor: The logits after applying the z-loss.
@@ -426,7 +480,7 @@ def apply_z_loss(self, logits):
         if self.config.moe_z_loss_coeff is not None and self.training and torch.is_grad_enabled():
             # Skip Z loss calculations when using torch.no_grad() or checkpointing.
             moe_z_loss_coeff = self.config.moe_z_loss_coeff / self.tp_cp_group.size()
-            z_loss = z_loss_func(logits, moe_z_loss_coeff)
+            z_loss = z_loss_func(logits, moe_z_loss_coeff, padding_mask=padding_mask)
             scale_up = 1.0
             if self.calculate_per_token_loss:
                 # The expected final scaling for z_loss gradients is
@@ -436,7 +490,8 @@ def apply_z_loss(self, logits):
                 # which scales both the main_loss gradient and z_loss gradient by
                 # 1/(num_local_tokens * dp_size * num_micro_batches) in finalize_model_grads().
                 # To correct this scaling, we need to scale the z_loss by num_local_tokens here.
-                logits = MoEAuxLossAutoScaler.apply(logits, z_loss * logits.shape[0])
+                num_tokens = padding_mask.sum() if padding_mask is not None else logits.shape[0]
+                logits = MoEAuxLossAutoScaler.apply(logits, z_loss * num_tokens)
             else:
                 logits = MoEAuxLossAutoScaler.apply(logits, z_loss)
 
@@ -470,20 +525,35 @@ def apply_input_jitter(self, input: torch.Tensor):
             return input
 
     @jit_fuser
-    def _apply_expert_bias(self, routing_map: torch.Tensor):
+    def _apply_expert_bias(self, routing_map: torch.Tensor, padding_mask: Optional[torch.Tensor] = None):
         """
         Update expert bias and tokens_per_expert
         Prevent extra local tokens accumulation on evaluation or activation recomputation
+        
+        Args:
+            routing_map (torch.Tensor): Token to expert routing map, shape [num_tokens, num_experts].
+            padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
+                Shape [num_tokens]. True for valid tokens, False for padding tokens.
         """
         if self.enable_expert_bias and torch.is_grad_enabled():
             with torch.no_grad():
-                self.local_tokens_per_expert += routing_map.sum(dim=0)
-
-    def routing(self, logits: torch.Tensor):
+                if padding_mask is not None:
+                    # Only count valid tokens, exclude padding tokens
+                    # This logic is consistent with compute_tokens_per_expert_with_padding
+                    mask_expanded = padding_mask.unsqueeze(-1)
+                    routing_map_masked = routing_map & mask_expanded
+                    self.local_tokens_per_expert += routing_map_masked.sum(dim=0)
+                else:
+                    self.local_tokens_per_expert += routing_map.sum(dim=0)
+
+    def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = None):
         """Top-k routing function
 
         Args:
             logits (torch.Tensor): Logits tensor after gating.
+            padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
+                                                   Shape in [seq_length, bsz]. True for valid tokens,
+                                                   False for padding tokens. Defaults to None.
 
         Returns:
             probs (torch.Tensor): The probabilities of token to experts assignment.
@@ -492,9 +562,15 @@ def routing(self, logits: torch.Tensor):
         """
         seq_length, bsz = logits.shape[:2]
         logits = logits.view(-1, self.config.num_moe_experts)
+        
+        # Flatten padding_mask to [num_tokens] if provided
+        if padding_mask is not None:
+            padding_mask_flat = padding_mask.reshape(-1)
+        else:
+            padding_mask_flat = None
 
         # Apply Z-Loss
-        logits = self.apply_z_loss(logits)
+        logits = self.apply_z_loss(logits, padding_mask=padding_mask_flat)
 
         # Calculate probs and routing_map for token dispatching
         if self.routing_type == "sinkhorn":
@@ -527,18 +603,22 @@ def routing(self, logits: torch.Tensor):
         if self.training and torch.is_grad_enabled() and self.is_aux_loss_enabled():
             # Calculate scores and routing_map for aux loss
             routing_map_for_aux_loss, scores_for_aux_loss = compute_routing_scores_for_aux_loss(
-                logits, self.topk, self.score_function, fused=self.config.moe_router_fusion
+                logits, self.topk, self.score_function, fused=self.config.moe_router_fusion,
+                padding_mask=padding_mask_flat
             )
-            probs = self._apply_aux_loss(probs, scores_for_aux_loss, routing_map_for_aux_loss)
+            probs = self._apply_aux_loss(probs, scores_for_aux_loss, routing_map_for_aux_loss,
+                                        padding_mask=padding_mask_flat)
             probs = self._apply_seq_aux_loss(
-                probs, scores_for_aux_loss, routing_map_for_aux_loss, seq_length, bsz
+                probs, scores_for_aux_loss, routing_map_for_aux_loss, seq_length, bsz,
+                padding_mask=padding_mask_flat
             )
             probs = self._apply_global_aux_loss(
-                probs, scores_for_aux_loss, routing_map_for_aux_loss
+                probs, scores_for_aux_loss, routing_map_for_aux_loss,
+                padding_mask=padding_mask_flat
             )
 
         # Optionally apply expert bias
-        self._apply_expert_bias(routing_map)
+        self._apply_expert_bias(routing_map, padding_mask=padding_mask_flat)
 
         return probs, routing_map
 
@@ -548,12 +628,15 @@ def reset_global_aux_loss_tracker(self):
             self.global_tokens_per_expert.zero_()
             self.ga_steps.zero_()
 
-    def forward(self, input: torch.Tensor):
+    def forward(self, input: torch.Tensor, padding_mask: Optional[torch.Tensor] = None):
         """
         Forward pass of the router.
 
         Args:
             input (torch.Tensor): Input tensor.
+            padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
+                                                   Shape in [seq_length, bsz]. True for valid tokens,
+                                                   False for padding tokens. Defaults to None.
         """
         self._maintain_float32_expert_bias()
 
@@ -565,7 +648,7 @@ def forward(self, input: torch.Tensor):
             # Apply force load balancing with random logits for benchmark
             logits = apply_random_logits(logits)
 
-        probs, routing_map = self.routing(logits)
+        probs, routing_map = self.routing(logits, padding_mask=padding_mask)
 
         return probs, routing_map
 
diff --git a/megatron/core/transformer/transformer_block.py b/megatron/core/transformer/transformer_block.py
index f189d5fa7a4..5ca13ff212e 100755
--- a/megatron/core/transformer/transformer_block.py
+++ b/megatron/core/transformer/transformer_block.py
@@ -389,7 +389,6 @@ def build_layer(layer_spec, layer_number):
     def has_final_layernorm_in_this_stage(self):
         """
         Check if this vpp stage contains the final layernorm.
-
         Note:
             Final layernorm now has been moved from the post-process stage to the last decoder
             layer by using this function.
@@ -598,6 +597,7 @@ def forward(
         inference_context: Optional[BaseInferenceContext] = None,
         packed_seq_params: Optional[PackedSeqParams] = None,
         sequence_len_offset: Optional[Tensor] = None,
+        padding_mask: Optional[Tensor] = None,
         *,
         inference_params: Optional[BaseInferenceContext] = None,
         dynamic_inference_decode_only: Optional[bool] = None,
@@ -707,6 +707,7 @@ def forward(
                     attention_bias=attention_bias,
                     packed_seq_params=packed_seq_params,
                     use_inner_quantization_context=use_inner_quantization_context,
+                    padding_mask=padding_mask,
                 )
             else:
                 for l_no, layer in enumerate(self.layers):
@@ -744,6 +745,7 @@ def forward(
                             inference_context=inference_context,
                             packed_seq_params=packed_seq_params,
                             sequence_len_offset=sequence_len_offset,
+                            padding_mask=padding_mask,
                         )
 
                     if (
diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py
index cacfb9d01b8..6be3eb3ed53 100644
--- a/megatron/core/transformer/transformer_layer.py
+++ b/megatron/core/transformer/transformer_layer.py
@@ -453,7 +453,12 @@ def forward(self, *args, **kwargs):
         # runners in the cuda graph manager
         kwargs.pop("dynamic_inference_decode_only", None)
         hidden_states, context = self._forward_attention(*args, **kwargs)
-        output = self._forward_mlp(hidden_states, kwargs.get("inference_context", None))
+        
+        output = self._forward_mlp(
+            hidden_states,
+            kwargs.get("inference_context", None),
+            padding_mask=kwargs.get("padding_mask", None),
+        )
         return output, context
 
     def _forward_attention(
@@ -470,6 +475,7 @@ def _forward_attention(
         inference_context: Optional[Any] = None,
         packed_seq_params: Optional[PackedSeqParams] = None,
         sequence_len_offset: Optional[Tensor] = None,
+        padding_mask: Optional[Tensor] = None,
         *,
         inference_params: Optional[Any] = None,
     ):
@@ -587,12 +593,14 @@ def _forward_attention(
 
         return hidden_states, context
 
-    def _forward_mlp(self, hidden_states, inference_context=None):
+    def _forward_mlp(self, hidden_states, inference_context=None, padding_mask=None):
         """
         Perform a forward pass through the feed-forward layer.
 
         Args:
             hidden_states (Tensor): Transformed hidden states before the MLP layernorm.
+            inference_context: Inference context for optimizations.
+            padding_mask (Tensor, optional): Padding mask for MoE routing. Shape [bsz, seq_length].
 
         Returns:
             output (Tensor): Transformed hidden states of shape [s, b, h].
@@ -605,6 +613,12 @@ def _forward_mlp(self, hidden_states, inference_context=None):
 
         # Residual connection.
         residual = hidden_states
+        
+        padding_mask_for_moe = None
+        if self.is_moe_layer and padding_mask is not None:
+            # padding_mask from GPTModel: [bsz, seq_length]
+            # MoE layer needs: [seq_length, bsz]
+            padding_mask_for_moe = padding_mask.transpose(0, 1).bool()
 
         if self.offload_mlp_norm:
             hidden_states = fine_grained_offloading_group_start(hidden_states, name="mlp_norm")
@@ -638,24 +652,46 @@ def _forward_mlp(self, hidden_states, inference_context=None):
             assert (
                 not self.recompute_pre_mlp_layernorm
             ), "Recomputation is not supported for CUDA graph."
-            cudagraph_outputs = self.mlp(pre_mlp_layernorm_output)
+            cudagraph_outputs = self.mlp(pre_mlp_layernorm_output, padding_mask=padding_mask_for_moe)
             return cudagraph_outputs + [residual]
         elif self.recompute_mlp:
             if self.config.fp8:
                 # import here to avoid circular import
                 from megatron.core.extensions.transformer_engine import te_checkpoint
 
-                mlp_output_with_bias = te_checkpoint(
-                    self.mlp,
-                    False,
-                    tensor_parallel.random.get_cuda_rng_tracker,
-                    self.pg_collection.tp,
-                    pre_mlp_layernorm_output,
-                )
+                if self.is_moe_layer:
+                    # For MoE with checkpointing, we need a wrapper to pass padding_mask
+                    def mlp_forward_with_padding(hidden_states):
+                        return self.mlp(hidden_states, padding_mask=padding_mask_for_moe)
+                    
+                    mlp_output_with_bias = te_checkpoint(
+                        mlp_forward_with_padding,
+                        False,
+                        tensor_parallel.random.get_cuda_rng_tracker,
+                        self.pg_collection.tp,
+                        pre_mlp_layernorm_output,
+                    )
+                else:
+                    mlp_output_with_bias = te_checkpoint(
+                        self.mlp,
+                        False,
+                        tensor_parallel.random.get_cuda_rng_tracker,
+                        self.pg_collection.tp,
+                        pre_mlp_layernorm_output,
+                    )
             else:
-                mlp_output_with_bias = tensor_parallel.checkpoint(
-                    self.mlp, False, pre_mlp_layernorm_output
-                )
+                if self.is_moe_layer:
+                    # For MoE with checkpointing, we need a wrapper to pass padding_mask
+                    def mlp_forward_with_padding(hidden_states):
+                        return self.mlp(hidden_states, padding_mask=padding_mask_for_moe)
+                    
+                    mlp_output_with_bias = tensor_parallel.checkpoint(
+                        mlp_forward_with_padding, False, pre_mlp_layernorm_output
+                    )
+                else:
+                    mlp_output_with_bias = tensor_parallel.checkpoint(
+                        self.mlp, False, pre_mlp_layernorm_output
+                    )
         elif should_chunk_mlp_for_prefill:
             # Chunk input along sequence dimension
             num_chunks = min(self.config.mlp_chunks_for_prefill, pre_mlp_layernorm_output.shape[0])
@@ -670,7 +706,10 @@ def _forward_mlp(self, hidden_states, inference_context=None):
             bias_output = torch.stack(bias_chunks, dim=0).sum(dim=0) if bias_chunks else None
             mlp_output_with_bias = (mlp_output, bias_output)
         else:
-            mlp_output_with_bias = self.mlp(pre_mlp_layernorm_output)
+            if self.is_moe_layer:
+                mlp_output_with_bias = self.mlp(pre_mlp_layernorm_output, padding_mask=padding_mask_for_moe)
+            else:
+                mlp_output_with_bias = self.mlp(pre_mlp_layernorm_output)
 
         if self.recompute_pre_mlp_layernorm:
             # discard the output of the pre-mlp layernorm and register the recompute
diff --git a/tests/unit_tests/transformer/moe/test_aux_loss.py b/tests/unit_tests/transformer/moe/test_aux_loss.py
index 108c9a65a4b..e7dfb65c704 100644
--- a/tests/unit_tests/transformer/moe/test_aux_loss.py
+++ b/tests/unit_tests/transformer/moe/test_aux_loss.py
@@ -575,3 +575,145 @@ def test_force_balanced_aux_loss(self, tp_size, ep_size, cp_size):
                 reduce_from_tensor_model_parallel_region(aux_loss, router.tp_cp_group)
             assert aux_loss.item() == 1, f"{aux_loss_type}: {aux_loss.item()}"
             clear_aux_losses_tracker()
+
+
+class TestPaddingMaskAuxLoss:
+    """Test padding mask support in various aux loss types."""
+    
+    def setup_method(self, method):
+        Utils.initialize_model_parallel(1, 1)
+        _set_random_seed(seed_=123, data_parallel_random_init=False)
+
+        # Default configuration
+        self.default_transformer_config = TransformerConfig(
+            num_layers=1,
+            hidden_size=12,
+            num_attention_heads=8,
+            num_moe_experts=32,
+            use_cpu_initialization=True,
+            moe_router_load_balancing_type="aux_loss",
+            moe_router_topk=8,
+            moe_aux_loss_coeff=1.0,
+            bf16=True,
+            params_dtype=torch.bfloat16,
+            add_bias_linear=False,
+        )
+
+    def new_router(self, **kwargs):
+        """Create a new router with updated configuration."""
+        pg_collection = get_default_pg_collection()
+        new_transformer_config = dataclasses.replace(self.default_transformer_config, **kwargs)
+        router = TopKRouter(config=new_transformer_config, pg_collection=pg_collection)
+        router.set_layer_number(0)
+        return router
+
+    def teardown_method(self, method):
+        Utils.destroy_model_parallel()
+
+    @pytest.mark.internal
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+    @pytest.mark.parametrize("aux_loss_type", ["aux_loss", "seq_aux_loss", "global_aux_loss"])
+    def test_padding_mask_removes_padding_tokens(self, aux_loss_type):
+        """Test that padding tokens are correctly excluded from aux loss calculation."""
+        clear_aux_losses_tracker()
+        
+        router = self.new_router(
+            moe_router_load_balancing_type=aux_loss_type,
+            moe_aux_loss_coeff=1.0,
+            moe_router_dtype="fp64",
+        ).cuda()
+        
+        seq_len = 32
+        batch_size = 2
+        hidden_size = router.config.hidden_size
+        
+        # Create input with padding
+        hidden_states_full = torch.randn(
+            (seq_len, batch_size, hidden_size), dtype=torch.bfloat16, device='cuda'
+        )
+        
+        # Create padding mask: first half valid, second half padding
+        padding_mask = torch.ones((seq_len, batch_size), dtype=torch.bool, device='cuda')
+        padding_mask[seq_len // 2:, :] = False
+        
+        # Test with padding mask
+        router.weight.grad = None
+        scores_with_mask, routing_map_with_mask = router(hidden_states_full, padding_mask=padding_mask)
+        scores_with_mask.backward(torch.zeros_like(scores_with_mask))
+        
+        loss_name = {
+            "aux_loss": "load_balancing_loss",
+            "seq_aux_loss": "seq_load_balancing_loss",
+            "global_aux_loss": "global_load_balancing_loss",
+        }[aux_loss_type]
+        
+        tracker = get_moe_layer_wise_logging_tracker()
+        aux_loss_with_mask = tracker[loss_name]["values"][0].clone()
+        grad_with_mask = router.weight.grad.clone()
+        
+        # Test without padding (with only half of the tokens)
+        clear_aux_losses_tracker()
+        router.weight.grad = None
+        hidden_states_valid = hidden_states_full[:seq_len // 2, :, :]
+        scores_without_mask, routing_map_without_mask = router(hidden_states_valid)
+        scores_without_mask.backward(torch.zeros_like(scores_without_mask))
+        
+        aux_loss_without_mask = tracker[loss_name]["values"][0].clone()
+        grad_without_mask = router.weight.grad.clone()
+        
+        # The aux loss with mask should be close to the aux loss without mask
+        torch.testing.assert_close(aux_loss_with_mask, aux_loss_without_mask, rtol=1e-2, atol=1e-3)
+        torch.testing.assert_close(grad_with_mask, grad_without_mask, rtol=1e-2, atol=1e-3)
+        
+        clear_aux_losses_tracker()
+
+    @pytest.mark.internal
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+    def test_padding_mask_with_z_loss(self):
+        """Test that padding mask works correctly with z_loss."""
+        clear_aux_losses_tracker()
+        
+        router = self.new_router(
+            moe_router_load_balancing_type="aux_loss",
+            moe_aux_loss_coeff=0.0,
+            moe_z_loss_coeff=1.0,
+            moe_router_dtype="fp32",
+        ).cuda()
+        
+        seq_len = 32
+        batch_size = 2
+        hidden_size = router.config.hidden_size
+        
+        # Create input
+        hidden_states_full = torch.randn(
+            (seq_len, batch_size, hidden_size), dtype=torch.bfloat16, device='cuda'
+        )
+        
+        # Create padding mask: first half valid, second half padding
+        padding_mask = torch.ones((seq_len, batch_size), dtype=torch.bool, device='cuda')
+        padding_mask[seq_len // 2:, :] = False
+        
+        # Test with padding mask
+        router.weight.grad = None
+        scores_with_mask, _ = router(hidden_states_full, padding_mask=padding_mask)
+        scores_with_mask.sum().backward()
+        
+        tracker = get_moe_layer_wise_logging_tracker()
+        z_loss_with_mask = tracker["z_loss"]["values"][0].clone()
+        grad_with_mask = router.weight.grad.clone()
+        
+        # Test without padding (with only half of the tokens)
+        clear_aux_losses_tracker()
+        router.weight.grad = None
+        hidden_states_valid = hidden_states_full[:seq_len // 2, :, :]
+        scores_without_mask, _ = router(hidden_states_valid)
+        scores_without_mask.sum().backward()
+        
+        z_loss_without_mask = tracker["z_loss"]["values"][0].clone()
+        grad_without_mask = router.weight.grad.clone()
+        
+        # The z_loss with mask should be close to the z_loss without mask
+        torch.testing.assert_close(z_loss_with_mask, z_loss_without_mask, rtol=1e-2, atol=1e-3)
+        torch.testing.assert_close(grad_with_mask, grad_without_mask, rtol=1e-2, atol=1e-3)
+        
+        clear_aux_losses_tracker()
diff --git a/tests/unit_tests/transformer/moe/test_routers.py b/tests/unit_tests/transformer/moe/test_routers.py
index 677d938cdc7..e1198d96d5b 100644
--- a/tests/unit_tests/transformer/moe/test_routers.py
+++ b/tests/unit_tests/transformer/moe/test_routers.py
@@ -125,6 +125,42 @@ def test_aux_loss(self):
         out.sum().mul_(0).backward()
         assert self.sequential_mlp.router.weight.grad.abs().sum() > 0
 
+    @pytest.mark.internal
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+    def test_router_with_padding_mask(self):
+        """Test that padding mask correctly excludes padding tokens from routing."""
+        self.router = self.router.cuda()
+        seq_len = 32
+        batch_size = 2
+        hidden_size = self.router.config.hidden_size
+        
+        # Create input with shape [seq_len, batch_size, hidden_size]
+        hidden_states = torch.randn((seq_len, batch_size, hidden_size)).cuda().bfloat16()
+        
+        # Create padding mask: first half valid, second half padding
+        # padding_mask shape: [seq_len, batch_size]
+        padding_mask = torch.ones((seq_len, batch_size), dtype=torch.bool, device='cuda')
+        padding_mask[seq_len // 2:, :] = False  # Second half is padding
+        
+        # Test forward pass with padding mask
+        with torch.no_grad():
+            probs_with_mask, routing_map_with_mask = self.router(hidden_states, padding_mask=padding_mask)
+            
+            # Test forward pass without padding mask (only valid tokens)
+            hidden_states_valid = hidden_states[:seq_len // 2, :, :]
+            probs_without_mask, routing_map_without_mask = self.router(hidden_states_valid)
+            
+            # The valid part of routing with mask should match routing without mask
+            probs_valid_part = probs_with_mask.reshape(seq_len, batch_size, -1)[:seq_len // 2, :, :]
+            probs_valid_part = probs_valid_part.reshape(-1, probs_valid_part.shape[-1])
+            
+            # Check that shapes are as expected
+            assert probs_with_mask.shape == (seq_len * batch_size, self.router.config.num_moe_experts)
+            assert routing_map_with_mask.shape == (seq_len * batch_size, self.router.config.num_moe_experts)
+            
+            # Verify that probs for valid tokens are similar
+            torch.testing.assert_close(probs_valid_part, probs_without_mask, rtol=1e-3, atol=1e-3)
+
     @pytest.mark.internal
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
     def test_router_dtype(self):

From 5781e60f9ca301433259cd2720f2696295d9422a Mon Sep 17 00:00:00 2001
From: HaochenYuan <haocheny@nvidia.com>
Date: Thu, 6 Nov 2025 22:28:12 -0800
Subject: [PATCH 02/32] format fix

---
 megatron/core/transformer/moe/moe_layer.py | 12 +--
 megatron/core/transformer/moe/router.py    | 93 +++++++++++++---------
 2 files changed, 63 insertions(+), 42 deletions(-)

diff --git a/megatron/core/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py
index 5c8ca3181ab..3a93b9709aa 100644
--- a/megatron/core/transformer/moe/moe_layer.py
+++ b/megatron/core/transformer/moe/moe_layer.py
@@ -183,11 +183,11 @@ def route(self, hidden_states: torch.Tensor, padding_mask: Optional[torch.Tensor
 
         This method uses the router to determine which experts to send each token to,
         producing routing probabilities and a mapping.
-        
+
         Args:
             hidden_states (torch.Tensor): Input tensor.
             padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
-                                                   Shape in [seq_length, bsz]. True for valid tokens,
+                                                   Shape = [seq_length, bsz]. True for valid tokens,
                                                    False for padding tokens. Defaults to None.
         """
         probs, routing_map = self.router(hidden_states, padding_mask=padding_mask)
@@ -286,9 +286,9 @@ def forward(self, hidden_states: torch.Tensor, padding_mask: Optional[torch.Tens
         4. Combine: The outputs from the experts are combined and returned.
 
         Args:
-            hidden_states (torch.Tensor): The input tensor to the MoE layer. 
+            hidden_states (torch.Tensor): The input tensor to the MoE layer.
             padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
-                                                   Shape in [seq_length, bsz]. True for valid tokens,
+                                                   Shape = [seq_length, bsz]. True for valid tokens,
                                                    False for padding tokens. Defaults to None.
 
         Returns:
@@ -330,7 +330,9 @@ def custom_forward(hidden_states, padding_mask=None):
                     padding_mask,
                 )
             else:
-                outputs = tensor_parallel.checkpoint(custom_forward, False, hidden_states, padding_mask)
+                outputs = tensor_parallel.checkpoint(
+                    custom_forward, False, hidden_states, padding_mask
+                )
         else:
             outputs = custom_forward(hidden_states, padding_mask)
 
diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py
index 83d0f8c131e..c84713fdb0f 100644
--- a/megatron/core/transformer/moe/router.py
+++ b/megatron/core/transformer/moe/router.py
@@ -269,21 +269,22 @@ def is_aux_loss_enabled(self) -> bool:
         return False
 
     def _apply_aux_loss(
-        self, probs: torch.Tensor, scores_for_aux_loss: torch.Tensor, routing_map: torch.Tensor,
-        padding_mask: Optional[torch.Tensor] = None
+        self,
+        probs: torch.Tensor,
+        scores_for_aux_loss: torch.Tensor,
+        routing_map: torch.Tensor,
+        padding_mask: Optional[torch.Tensor] = None,
     ):
         """Apply the auxiliary loss for the given scores and routing map."""
         aux_loss_coeff = self.get_aux_loss_coeff("aux_loss")
         if aux_loss_coeff == 0:
             return probs
-        
+
         # Use unified function to compute tokens_per_expert and num_tokens
         tokens_per_expert, num_tokens = compute_tokens_per_expert_with_padding(
-            routing_map=routing_map,
-            padding_mask=padding_mask,
-            reshape_for_seq_aux=False,
+            routing_map=routing_map, padding_mask=padding_mask, reshape_for_seq_aux=False
         )
-        
+
         tokens_per_expert = reduce_from_tensor_model_parallel_region(
             tokens_per_expert, self.tp_cp_group
         )
@@ -300,8 +301,12 @@ def _apply_aux_loss(
             padding_mask=padding_mask,
         )
         probs = self.attach_and_log_load_balancing_loss(
-            probs, aux_loss_coeff, aux_loss, "load_balancing_loss", self.tp_cp_group,
-            valid_token_count=num_tokens
+            probs,
+            aux_loss_coeff,
+            aux_loss,
+            "load_balancing_loss",
+            self.tp_cp_group,
+            valid_token_count=num_tokens,
         )
         return probs
 
@@ -326,7 +331,7 @@ def _apply_seq_aux_loss(
             return probs
 
         scores_for_aux_loss = scores_for_aux_loss.reshape(seq_length, -1)
-        
+
         # Use unified function to compute tokens_per_expert and num_tokens
         tokens_per_expert, num_tokens = compute_tokens_per_expert_with_padding(
             routing_map=routing_map,
@@ -336,9 +341,9 @@ def _apply_seq_aux_loss(
             bsz=bsz,
             num_experts=self.config.num_moe_experts,
         )
-        
+
         total_num_tokens = num_tokens * self.tp_cp_group.size()
-        
+
         tokens_per_expert = reduce_from_tensor_model_parallel_region(
             tokens_per_expert, self.tp_cp_group
         )
@@ -361,16 +366,23 @@ def _apply_seq_aux_loss(
             valid_token_count = padding_mask.sum()
         else:
             valid_token_count = seq_length * bsz
-        
+
         probs = self.attach_and_log_load_balancing_loss(
-            probs, seq_aux_loss_coeff, aux_loss, "seq_load_balancing_loss", self.tp_cp_group,
-            valid_token_count=valid_token_count
+            probs,
+            seq_aux_loss_coeff,
+            aux_loss,
+            "seq_load_balancing_loss",
+            self.tp_cp_group,
+            valid_token_count=valid_token_count,
         )
         return probs
 
     def _apply_global_aux_loss(
-        self, probs: torch.Tensor, scores_for_aux_loss: torch.Tensor, routing_map: torch.Tensor,
-        padding_mask: Optional[torch.Tensor] = None
+        self,
+        probs: torch.Tensor,
+        scores_for_aux_loss: torch.Tensor,
+        routing_map: torch.Tensor,
+        padding_mask: Optional[torch.Tensor] = None,
     ):
         """Apply the global auxiliary loss for the given scores and routing map."""
         global_aux_loss_coeff = self.get_aux_loss_coeff("global_aux_loss")
@@ -379,11 +391,9 @@ def _apply_global_aux_loss(
 
         # Use unified function to compute tokens_per_expert and num_tokens
         tokens_per_expert, num_tokens = compute_tokens_per_expert_with_padding(
-            routing_map=routing_map,
-            padding_mask=padding_mask,
-            reshape_for_seq_aux=False,
+            routing_map=routing_map, padding_mask=padding_mask, reshape_for_seq_aux=False
         )
-        
+
         tokens_per_expert = reduce_from_tensor_model_parallel_region(
             tokens_per_expert, self.tp_dp_cp_group
         )
@@ -424,14 +434,14 @@ def attach_and_log_load_balancing_loss(
         valid_token_count: Optional[Union[int, torch.Tensor]] = None,
     ):
         """Attach aux loss function to activation and add to logging.
-        
+
         Args:
             activation (torch.Tensor): Activation tensor to attach the aux loss to.
             aux_loss_coeff (float): Coefficient for the aux loss.
             aux_loss (torch.Tensor): Computed aux loss.
             aux_loss_name (str): Name of the aux loss for logging.
             reduce_group (torch.distributed.ProcessGroup): Process group for reduction.
-            valid_token_count (int or torch.Tensor, optional): Number of valid tokens excluding 
+            valid_token_count (int or torch.Tensor, optional): Number of valid tokens excluding
                 padding tokens. Can be a Python int or a torch.Tensor (typically 0-d tensor).
                 If None, uses activation.shape[0]. Defaults to None.
         """
@@ -525,13 +535,15 @@ def apply_input_jitter(self, input: torch.Tensor):
             return input
 
     @jit_fuser
-    def _apply_expert_bias(self, routing_map: torch.Tensor, padding_mask: Optional[torch.Tensor] = None):
+    def _apply_expert_bias(
+        self, routing_map: torch.Tensor, padding_mask: Optional[torch.Tensor] = None
+    ):
         """
         Update expert bias and tokens_per_expert
         Prevent extra local tokens accumulation on evaluation or activation recomputation
-        
+
         Args:
-            routing_map (torch.Tensor): Token to expert routing map, shape [num_tokens, num_experts].
+            routing_map (torch.Tensor): Token to expert routing map, [num_tokens, num_experts].
             padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
                 Shape [num_tokens]. True for valid tokens, False for padding tokens.
         """
@@ -552,7 +564,7 @@ def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = N
         Args:
             logits (torch.Tensor): Logits tensor after gating.
             padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
-                                                   Shape in [seq_length, bsz]. True for valid tokens,
+                                                   Shape = [seq_length, bsz]. True for valid tokens,
                                                    False for padding tokens. Defaults to None.
 
         Returns:
@@ -562,7 +574,7 @@ def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = N
         """
         seq_length, bsz = logits.shape[:2]
         logits = logits.view(-1, self.config.num_moe_experts)
-        
+
         # Flatten padding_mask to [num_tokens] if provided
         if padding_mask is not None:
             padding_mask_flat = padding_mask.reshape(-1)
@@ -603,18 +615,25 @@ def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = N
         if self.training and torch.is_grad_enabled() and self.is_aux_loss_enabled():
             # Calculate scores and routing_map for aux loss
             routing_map_for_aux_loss, scores_for_aux_loss = compute_routing_scores_for_aux_loss(
-                logits, self.topk, self.score_function, fused=self.config.moe_router_fusion,
-                padding_mask=padding_mask_flat
+                logits,
+                self.topk,
+                self.score_function,
+                fused=self.config.moe_router_fusion,
+                padding_mask=padding_mask_flat,
+            )
+            probs = self._apply_aux_loss(
+                probs, scores_for_aux_loss, routing_map_for_aux_loss, padding_mask=padding_mask_flat
             )
-            probs = self._apply_aux_loss(probs, scores_for_aux_loss, routing_map_for_aux_loss,
-                                        padding_mask=padding_mask_flat)
             probs = self._apply_seq_aux_loss(
-                probs, scores_for_aux_loss, routing_map_for_aux_loss, seq_length, bsz,
-                padding_mask=padding_mask_flat
+                probs,
+                scores_for_aux_loss,
+                routing_map_for_aux_loss,
+                seq_length,
+                bsz,
+                padding_mask=padding_mask_flat,
             )
             probs = self._apply_global_aux_loss(
-                probs, scores_for_aux_loss, routing_map_for_aux_loss,
-                padding_mask=padding_mask_flat
+                probs, scores_for_aux_loss, routing_map_for_aux_loss, padding_mask=padding_mask_flat
             )
 
         # Optionally apply expert bias
@@ -635,7 +654,7 @@ def forward(self, input: torch.Tensor, padding_mask: Optional[torch.Tensor] = No
         Args:
             input (torch.Tensor): Input tensor.
             padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
-                                                   Shape in [seq_length, bsz]. True for valid tokens,
+                                                   Shape = [seq_length, bsz]. True for valid tokens,
                                                    False for padding tokens. Defaults to None.
         """
         self._maintain_float32_expert_bias()

From 7e4a3f256d63275a1bf6eac187f2cb899a1772a4 Mon Sep 17 00:00:00 2001
From: HaochenYuan <haocheny@nvidia.com>
Date: Thu, 6 Nov 2025 23:23:25 -0800
Subject: [PATCH 03/32] format fix

---
 megatron/core/transformer/moe/moe_utils.py    | 40 +++++++------
 .../core/transformer/transformer_layer.py     | 16 ++++--
 .../transformer/moe/test_aux_loss.py          | 56 ++++++++++---------
 .../transformer/moe/test_routers.py           | 36 +++++++-----
 4 files changed, 86 insertions(+), 62 deletions(-)

diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py
index 5f5c59f5b2a..ef435967c8d 100644
--- a/megatron/core/transformer/moe/moe_utils.py
+++ b/megatron/core/transformer/moe/moe_utils.py
@@ -44,7 +44,7 @@ def compute_tokens_per_expert_with_padding(
 ):
     """Compute tokens_per_expert and total_num_tokens with optional padding mask.
     This function provides a unified way to compute token counts across different aux loss types.
-    
+
     Args:
         routing_map (torch.Tensor): Token to expert routing map.
             - For aux_loss/global_aux_loss: shape [num_tokens, num_experts]
@@ -55,25 +55,28 @@ def compute_tokens_per_expert_with_padding(
         seq_length (int, optional): Sequence length, required when reshape_for_seq_aux=True.
         bsz (int, optional): Batch size, required when reshape_for_seq_aux=True.
         num_experts (int, optional): Number of experts, required when reshape_for_seq_aux=True.
-    
+
     Returns:
         tuple: (tokens_per_expert, num_valid_tokens)
-            - tokens_per_expert (torch.Tensor): Number of tokens per expert, shape [num_experts] 
+            - tokens_per_expert (torch.Tensor): Number of tokens per expert, shape [num_experts]
                 or [bsz * num_experts] for seq_aux_loss
             - num_valid_tokens (torch.Tensor or int): Number of valid (non-padding) tokens
     """
     if reshape_for_seq_aux:
         # seq aux loss
-        assert seq_length is not None and bsz is not None and num_experts is not None, \
-            "seq_length, bsz, and num_experts must be provided when reshape_for_seq_aux=True"
-        
+        assert (
+            seq_length is not None and bsz is not None and num_experts is not None
+        ), "seq_length, bsz, and num_experts must be provided when reshape_for_seq_aux=True"
+
         if padding_mask is not None:
             # Reshape padding_mask to [seq_length, bsz]
             padding_mask_reshaped = padding_mask.reshape(seq_length, bsz)
             # Expand to match routing_map after reshape [seq_length, bsz * num_experts]
-            mask_expanded = padding_mask_reshaped.unsqueeze(-1).expand(
-                -1, -1, num_experts
-            ).reshape(seq_length, -1)
+            mask_expanded = (
+                padding_mask_reshaped.unsqueeze(-1)
+                .expand(-1, -1, num_experts)
+                .reshape(seq_length, -1)
+            )
             routing_map_masked = routing_map.reshape(seq_length, -1) & mask_expanded
             tokens_per_expert = routing_map_masked.sum(dim=0)
             # Count valid tokens only
@@ -93,7 +96,7 @@ def compute_tokens_per_expert_with_padding(
         else:
             tokens_per_expert = routing_map.sum(dim=0)
             num_valid_tokens = routing_map.shape[0]
-    
+
     return tokens_per_expert, num_valid_tokens
 
 
@@ -167,7 +170,7 @@ def switch_load_balancing_loss_func(
         # padding_mask: [num_tokens], probs: [num_tokens, num_experts]
         mask_expanded = padding_mask.unsqueeze(-1)
         probs = probs * mask_expanded
-    
+
     if fused:
         if not HAVE_TE or fused_moe_aux_loss is None:
             raise ValueError("fused_moe_aux_loss is not available. Please install TE >= 2.7.0.")
@@ -203,7 +206,7 @@ def z_loss_func(logits, z_loss_coeff, padding_mask: Optional[torch.Tensor] = Non
     """
     logsum = torch.logsumexp(logits, dim=-1)
     z_loss_values = torch.square(logsum)
-    
+
     if padding_mask is not None:
         # Only compute z_loss for non-padding tokens
         z_loss_values = z_loss_values * padding_mask
@@ -710,8 +713,11 @@ def compute_topk(scores, topk, num_groups=None, group_topk=None):
 
 
 def compute_routing_scores_for_aux_loss(
-    logits: torch.Tensor, topk: int, score_function: str, fused: bool = False, 
-    padding_mask: Optional[torch.Tensor] = None
+    logits: torch.Tensor,
+    topk: int,
+    score_function: str,
+    fused: bool = False,
+    padding_mask: Optional[torch.Tensor] = None,
 ):
     """Compute routing scores based on the score function.
 
@@ -730,7 +736,9 @@ def compute_routing_scores_for_aux_loss(
             )
         # Note: fused implementation does not support padding_mask yet
         if padding_mask is not None:
-            raise ValueError("Fused compute_routing_scores does not support padding_mask. Set fused=False.")
+            raise ValueError(
+                "Fused compute_routing_scores does not support padding_mask. Set fused=False."
+            )
         return fused_compute_score_for_moe_aux_loss(
             logits=logits, topk=topk, score_function=score_function
         )
@@ -745,7 +753,7 @@ def compute_routing_scores_for_aux_loss(
 
     _, top_indices = torch.topk(scores, k=topk, dim=1)
     routing_map = torch.zeros_like(logits).int().scatter(1, top_indices, 1).bool()
-    
+
     # Apply padding mask to scores if provided
     if padding_mask is not None:
         mask_expanded = padding_mask.unsqueeze(-1).to(scores.dtype)
diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py
index 6be3eb3ed53..49eb7f2b9ec 100644
--- a/megatron/core/transformer/transformer_layer.py
+++ b/megatron/core/transformer/transformer_layer.py
@@ -453,7 +453,7 @@ def forward(self, *args, **kwargs):
         # runners in the cuda graph manager
         kwargs.pop("dynamic_inference_decode_only", None)
         hidden_states, context = self._forward_attention(*args, **kwargs)
-        
+
         output = self._forward_mlp(
             hidden_states,
             kwargs.get("inference_context", None),
@@ -613,7 +613,7 @@ def _forward_mlp(self, hidden_states, inference_context=None, padding_mask=None)
 
         # Residual connection.
         residual = hidden_states
-        
+
         padding_mask_for_moe = None
         if self.is_moe_layer and padding_mask is not None:
             # padding_mask from GPTModel: [bsz, seq_length]
@@ -652,7 +652,9 @@ def _forward_mlp(self, hidden_states, inference_context=None, padding_mask=None)
             assert (
                 not self.recompute_pre_mlp_layernorm
             ), "Recomputation is not supported for CUDA graph."
-            cudagraph_outputs = self.mlp(pre_mlp_layernorm_output, padding_mask=padding_mask_for_moe)
+            cudagraph_outputs = self.mlp(
+                pre_mlp_layernorm_output, padding_mask=padding_mask_for_moe
+            )
             return cudagraph_outputs + [residual]
         elif self.recompute_mlp:
             if self.config.fp8:
@@ -663,7 +665,7 @@ def _forward_mlp(self, hidden_states, inference_context=None, padding_mask=None)
                     # For MoE with checkpointing, we need a wrapper to pass padding_mask
                     def mlp_forward_with_padding(hidden_states):
                         return self.mlp(hidden_states, padding_mask=padding_mask_for_moe)
-                    
+
                     mlp_output_with_bias = te_checkpoint(
                         mlp_forward_with_padding,
                         False,
@@ -684,7 +686,7 @@ def mlp_forward_with_padding(hidden_states):
                     # For MoE with checkpointing, we need a wrapper to pass padding_mask
                     def mlp_forward_with_padding(hidden_states):
                         return self.mlp(hidden_states, padding_mask=padding_mask_for_moe)
-                    
+
                     mlp_output_with_bias = tensor_parallel.checkpoint(
                         mlp_forward_with_padding, False, pre_mlp_layernorm_output
                     )
@@ -707,7 +709,9 @@ def mlp_forward_with_padding(hidden_states):
             mlp_output_with_bias = (mlp_output, bias_output)
         else:
             if self.is_moe_layer:
-                mlp_output_with_bias = self.mlp(pre_mlp_layernorm_output, padding_mask=padding_mask_for_moe)
+                mlp_output_with_bias = self.mlp(
+                    pre_mlp_layernorm_output, padding_mask=padding_mask_for_moe
+                )
             else:
                 mlp_output_with_bias = self.mlp(pre_mlp_layernorm_output)
 
diff --git a/tests/unit_tests/transformer/moe/test_aux_loss.py b/tests/unit_tests/transformer/moe/test_aux_loss.py
index e7dfb65c704..83616c69970 100644
--- a/tests/unit_tests/transformer/moe/test_aux_loss.py
+++ b/tests/unit_tests/transformer/moe/test_aux_loss.py
@@ -579,7 +579,7 @@ def test_force_balanced_aux_loss(self, tp_size, ep_size, cp_size):
 
 class TestPaddingMaskAuxLoss:
     """Test padding mask support in various aux loss types."""
-    
+
     def setup_method(self, method):
         Utils.initialize_model_parallel(1, 1)
         _set_random_seed(seed_=123, data_parallel_random_init=False)
@@ -616,55 +616,57 @@ def teardown_method(self, method):
     def test_padding_mask_removes_padding_tokens(self, aux_loss_type):
         """Test that padding tokens are correctly excluded from aux loss calculation."""
         clear_aux_losses_tracker()
-        
+
         router = self.new_router(
             moe_router_load_balancing_type=aux_loss_type,
             moe_aux_loss_coeff=1.0,
             moe_router_dtype="fp64",
         ).cuda()
-        
+
         seq_len = 32
         batch_size = 2
         hidden_size = router.config.hidden_size
-        
+
         # Create input with padding
         hidden_states_full = torch.randn(
             (seq_len, batch_size, hidden_size), dtype=torch.bfloat16, device='cuda'
         )
-        
+
         # Create padding mask: first half valid, second half padding
         padding_mask = torch.ones((seq_len, batch_size), dtype=torch.bool, device='cuda')
-        padding_mask[seq_len // 2:, :] = False
-        
+        padding_mask[seq_len // 2 :, :] = False
+
         # Test with padding mask
         router.weight.grad = None
-        scores_with_mask, routing_map_with_mask = router(hidden_states_full, padding_mask=padding_mask)
+        scores_with_mask, routing_map_with_mask = router(
+            hidden_states_full, padding_mask=padding_mask
+        )
         scores_with_mask.backward(torch.zeros_like(scores_with_mask))
-        
+
         loss_name = {
             "aux_loss": "load_balancing_loss",
             "seq_aux_loss": "seq_load_balancing_loss",
             "global_aux_loss": "global_load_balancing_loss",
         }[aux_loss_type]
-        
+
         tracker = get_moe_layer_wise_logging_tracker()
         aux_loss_with_mask = tracker[loss_name]["values"][0].clone()
         grad_with_mask = router.weight.grad.clone()
-        
+
         # Test without padding (with only half of the tokens)
         clear_aux_losses_tracker()
         router.weight.grad = None
-        hidden_states_valid = hidden_states_full[:seq_len // 2, :, :]
+        hidden_states_valid = hidden_states_full[: seq_len // 2, :, :]
         scores_without_mask, routing_map_without_mask = router(hidden_states_valid)
         scores_without_mask.backward(torch.zeros_like(scores_without_mask))
-        
+
         aux_loss_without_mask = tracker[loss_name]["values"][0].clone()
         grad_without_mask = router.weight.grad.clone()
-        
+
         # The aux loss with mask should be close to the aux loss without mask
         torch.testing.assert_close(aux_loss_with_mask, aux_loss_without_mask, rtol=1e-2, atol=1e-3)
         torch.testing.assert_close(grad_with_mask, grad_without_mask, rtol=1e-2, atol=1e-3)
-        
+
         clear_aux_losses_tracker()
 
     @pytest.mark.internal
@@ -672,48 +674,48 @@ def test_padding_mask_removes_padding_tokens(self, aux_loss_type):
     def test_padding_mask_with_z_loss(self):
         """Test that padding mask works correctly with z_loss."""
         clear_aux_losses_tracker()
-        
+
         router = self.new_router(
             moe_router_load_balancing_type="aux_loss",
             moe_aux_loss_coeff=0.0,
             moe_z_loss_coeff=1.0,
             moe_router_dtype="fp32",
         ).cuda()
-        
+
         seq_len = 32
         batch_size = 2
         hidden_size = router.config.hidden_size
-        
+
         # Create input
         hidden_states_full = torch.randn(
             (seq_len, batch_size, hidden_size), dtype=torch.bfloat16, device='cuda'
         )
-        
+
         # Create padding mask: first half valid, second half padding
         padding_mask = torch.ones((seq_len, batch_size), dtype=torch.bool, device='cuda')
-        padding_mask[seq_len // 2:, :] = False
-        
+        padding_mask[seq_len // 2 :, :] = False
+
         # Test with padding mask
         router.weight.grad = None
         scores_with_mask, _ = router(hidden_states_full, padding_mask=padding_mask)
         scores_with_mask.sum().backward()
-        
+
         tracker = get_moe_layer_wise_logging_tracker()
         z_loss_with_mask = tracker["z_loss"]["values"][0].clone()
         grad_with_mask = router.weight.grad.clone()
-        
+
         # Test without padding (with only half of the tokens)
         clear_aux_losses_tracker()
         router.weight.grad = None
-        hidden_states_valid = hidden_states_full[:seq_len // 2, :, :]
+        hidden_states_valid = hidden_states_full[: seq_len // 2, :, :]
         scores_without_mask, _ = router(hidden_states_valid)
         scores_without_mask.sum().backward()
-        
+
         z_loss_without_mask = tracker["z_loss"]["values"][0].clone()
         grad_without_mask = router.weight.grad.clone()
-        
+
         # The z_loss with mask should be close to the z_loss without mask
         torch.testing.assert_close(z_loss_with_mask, z_loss_without_mask, rtol=1e-2, atol=1e-3)
         torch.testing.assert_close(grad_with_mask, grad_without_mask, rtol=1e-2, atol=1e-3)
-        
+
         clear_aux_losses_tracker()
diff --git a/tests/unit_tests/transformer/moe/test_routers.py b/tests/unit_tests/transformer/moe/test_routers.py
index e1198d96d5b..d11aef207c4 100644
--- a/tests/unit_tests/transformer/moe/test_routers.py
+++ b/tests/unit_tests/transformer/moe/test_routers.py
@@ -133,31 +133,41 @@ def test_router_with_padding_mask(self):
         seq_len = 32
         batch_size = 2
         hidden_size = self.router.config.hidden_size
-        
+
         # Create input with shape [seq_len, batch_size, hidden_size]
         hidden_states = torch.randn((seq_len, batch_size, hidden_size)).cuda().bfloat16()
-        
+
         # Create padding mask: first half valid, second half padding
         # padding_mask shape: [seq_len, batch_size]
         padding_mask = torch.ones((seq_len, batch_size), dtype=torch.bool, device='cuda')
-        padding_mask[seq_len // 2:, :] = False  # Second half is padding
-        
+        padding_mask[seq_len // 2 :, :] = False  # Second half is padding
+
         # Test forward pass with padding mask
         with torch.no_grad():
-            probs_with_mask, routing_map_with_mask = self.router(hidden_states, padding_mask=padding_mask)
-            
+            probs_with_mask, routing_map_with_mask = self.router(
+                hidden_states, padding_mask=padding_mask
+            )
+
             # Test forward pass without padding mask (only valid tokens)
-            hidden_states_valid = hidden_states[:seq_len // 2, :, :]
+            hidden_states_valid = hidden_states[: seq_len // 2, :, :]
             probs_without_mask, routing_map_without_mask = self.router(hidden_states_valid)
-            
+
             # The valid part of routing with mask should match routing without mask
-            probs_valid_part = probs_with_mask.reshape(seq_len, batch_size, -1)[:seq_len // 2, :, :]
+            probs_valid_part = probs_with_mask.reshape(seq_len, batch_size, -1)[
+                : seq_len // 2, :, :
+            ]
             probs_valid_part = probs_valid_part.reshape(-1, probs_valid_part.shape[-1])
-            
+
             # Check that shapes are as expected
-            assert probs_with_mask.shape == (seq_len * batch_size, self.router.config.num_moe_experts)
-            assert routing_map_with_mask.shape == (seq_len * batch_size, self.router.config.num_moe_experts)
-            
+            assert probs_with_mask.shape == (
+                seq_len * batch_size,
+                self.router.config.num_moe_experts,
+            )
+            assert routing_map_with_mask.shape == (
+                seq_len * batch_size,
+                self.router.config.num_moe_experts,
+            )
+
             # Verify that probs for valid tokens are similar
             torch.testing.assert_close(probs_valid_part, probs_without_mask, rtol=1e-3, atol=1e-3)
 

From f9d9c15580bbed7ec9dc8c86e8d95fe49b8fbaf7 Mon Sep 17 00:00:00 2001
From: HaochenYuan <haocheny@nvidia.com>
Date: Fri, 7 Nov 2025 03:28:24 -0800
Subject: [PATCH 04/32] fix UT error

---
 megatron/core/transformer/moe/router.py        | 5 +++--
 megatron/core/transformer/transformer_block.py | 6 +++++-
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py
index c84713fdb0f..8ff70f91f48 100644
--- a/megatron/core/transformer/moe/router.py
+++ b/megatron/core/transformer/moe/router.py
@@ -342,11 +342,12 @@ def _apply_seq_aux_loss(
             num_experts=self.config.num_moe_experts,
         )
 
-        total_num_tokens = num_tokens * self.tp_cp_group.size()
-
         tokens_per_expert = reduce_from_tensor_model_parallel_region(
             tokens_per_expert, self.tp_cp_group
         )
+        # For seq_aux_loss, total_num_tokens should be seq_length (per sequence),
+        # not seq_length * batch_size
+        total_num_tokens = seq_length * self.tp_cp_group.size()
 
         aux_loss = (
             switch_load_balancing_loss_func(
diff --git a/megatron/core/transformer/transformer_block.py b/megatron/core/transformer/transformer_block.py
index e5d3dd374a6..058c9d14b1e 100755
--- a/megatron/core/transformer/transformer_block.py
+++ b/megatron/core/transformer/transformer_block.py
@@ -427,12 +427,13 @@ def _checkpointed_forward(
         attention_bias: Tensor,
         packed_seq_params: PackedSeqParams,
         use_inner_quantization_context: bool,
+        padding_mask: Optional[Tensor] = None,
     ):
         """Forward method with activation checkpointing."""
 
         def custom(start: int, end: int):
             def custom_forward(
-                hidden_states, attention_mask, context, context_mask, rotary_pos_emb
+                hidden_states, attention_mask, context, context_mask, rotary_pos_emb, padding_mask
             ):
                 for index in range(start, end):
                     layer = self._get_layer(index)
@@ -463,6 +464,7 @@ def custom_forward(
                             attention_bias=attention_bias,
                             inference_context=None,
                             packed_seq_params=packed_seq_params,
+                            padding_mask=padding_mask,
                         )
                 return hidden_states, context
 
@@ -482,6 +484,7 @@ def checkpoint_handler(forward_func):
                     context,
                     context_mask,
                     rotary_pos_emb,
+                    padding_mask,
                 )
             else:
                 return tensor_parallel.checkpoint(
@@ -492,6 +495,7 @@ def checkpoint_handler(forward_func):
                     context,
                     context_mask,
                     rotary_pos_emb,
+                    padding_mask,
                 )
 
         if self.config.recompute_method == 'uniform':

From 2903bf9755f5bfbec7b84c42aa0ff5c75e2a9a4b Mon Sep 17 00:00:00 2001
From: HaochenYuan <haocheny@nvidia.com>
Date: Thu, 13 Nov 2025 00:56:15 -0800
Subject: [PATCH 05/32] fix some UT error

---
 megatron/core/transformer/moe/router.py       |  28 ++-
 .../core/transformer/transformer_block.py     |   2 +-
 .../transformer/moe/test_aux_loss.py          | 204 ++++++++++--------
 3 files changed, 144 insertions(+), 90 deletions(-)

diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py
index 8ff70f91f48..8a897db3ec3 100644
--- a/megatron/core/transformer/moe/router.py
+++ b/megatron/core/transformer/moe/router.py
@@ -288,7 +288,14 @@ def _apply_aux_loss(
         tokens_per_expert = reduce_from_tensor_model_parallel_region(
             tokens_per_expert, self.tp_cp_group
         )
-        total_num_tokens = num_tokens * self.tp_cp_group.size()
+        # When padding_mask is not None, total_num_tokens should be the sum across all ranks.
+        # After all_reduce, tokens_per_expert.sum() gives total token-expert assignments.
+        # Since each token is assigned to topk experts, we divide by topk to get actual token count.
+        # This avoids an extra all_reduce operation for num_tokens.
+        if padding_mask is not None:
+            total_num_tokens = tokens_per_expert.sum() / self.topk
+        else:
+            total_num_tokens = num_tokens * self.tp_cp_group.size()
 
         aux_loss = switch_load_balancing_loss_func(
             probs=scores_for_aux_loss,
@@ -346,8 +353,14 @@ def _apply_seq_aux_loss(
             tokens_per_expert, self.tp_cp_group
         )
         # For seq_aux_loss, total_num_tokens should be seq_length (per sequence),
-        # not seq_length * batch_size
-        total_num_tokens = seq_length * self.tp_cp_group.size()
+        # not seq_length * batch_size.
+        # When padding_mask is not None, compute from tokens_per_expert to get accurate count.
+        # tokens_per_expert shape: [bsz * num_experts], tokens_per_expert.sum() = seq_length * bsz * topk
+        # Divide by (bsz * topk) to get seq_length
+        if padding_mask is not None:
+            total_num_tokens = tokens_per_expert.sum() / (bsz * self.topk)
+        else:
+            total_num_tokens = seq_length * self.tp_cp_group.size()
 
         aux_loss = (
             switch_load_balancing_loss_func(
@@ -403,7 +416,14 @@ def _apply_global_aux_loss(
         self.ga_steps += 1
         averated_tokens_per_expert = self.global_tokens_per_expert / self.ga_steps
 
-        total_num_tokens = num_tokens * self.tp_dp_cp_group.size()
+        # When padding_mask is not None, total_num_tokens should be the sum across all ranks.
+        # After all_reduce, tokens_per_expert.sum() gives total token-expert assignments.
+        # Since each token is assigned to topk experts, we divide by topk to get actual token count.
+        # This avoids an extra all_reduce operation for num_tokens.
+        if padding_mask is not None:
+            total_num_tokens = tokens_per_expert.sum() / self.topk
+        else:
+            total_num_tokens = num_tokens * self.tp_dp_cp_group.size()
 
         global_aux_loss = switch_load_balancing_loss_func(
             probs=scores_for_aux_loss,
diff --git a/megatron/core/transformer/transformer_block.py b/megatron/core/transformer/transformer_block.py
index 058c9d14b1e..11c4512035d 100755
--- a/megatron/core/transformer/transformer_block.py
+++ b/megatron/core/transformer/transformer_block.py
@@ -433,7 +433,7 @@ def _checkpointed_forward(
 
         def custom(start: int, end: int):
             def custom_forward(
-                hidden_states, attention_mask, context, context_mask, rotary_pos_emb, padding_mask
+                hidden_states, attention_mask, context, context_mask, rotary_pos_emb, padding_mask=None
             ):
                 for index in range(start, end):
                     layer = self._get_layer(index)
diff --git a/tests/unit_tests/transformer/moe/test_aux_loss.py b/tests/unit_tests/transformer/moe/test_aux_loss.py
index 83616c69970..2cd1da6f381 100644
--- a/tests/unit_tests/transformer/moe/test_aux_loss.py
+++ b/tests/unit_tests/transformer/moe/test_aux_loss.py
@@ -580,10 +580,27 @@ def test_force_balanced_aux_loss(self, tp_size, ep_size, cp_size):
 class TestPaddingMaskAuxLoss:
     """Test padding mask support in various aux loss types."""
 
-    def setup_method(self, method):
-        Utils.initialize_model_parallel(1, 1)
+    def setup_model_parallel(self, tp_size=1, ep_size=1, cp_size=1):
+        """Initialize model parallel with given configuration.
+        
+        Args:
+            tp_size: Tensor parallel size.
+            ep_size: Expert parallel size.
+            cp_size: Context parallel size.
+        """
+        Utils.initialize_model_parallel(
+            tensor_model_parallel_size=tp_size,
+            pipeline_model_parallel_size=1,
+            context_parallel_size=cp_size,
+            expert_model_parallel_size=ep_size,
+        )
         _set_random_seed(seed_=123, data_parallel_random_init=False)
 
+        # Store parallel configuration
+        self.tp_size = tp_size
+        self.ep_size = ep_size
+        self.cp_size = cp_size
+
         # Default configuration
         self.default_transformer_config = TransformerConfig(
             num_layers=1,
@@ -607,115 +624,132 @@ def new_router(self, **kwargs):
         router.set_layer_number(0)
         return router
 
-    def teardown_method(self, method):
-        Utils.destroy_model_parallel()
-
     @pytest.mark.internal
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
     @pytest.mark.parametrize("aux_loss_type", ["aux_loss", "seq_aux_loss", "global_aux_loss"])
-    def test_padding_mask_removes_padding_tokens(self, aux_loss_type):
+    @pytest.mark.parametrize(
+        "tp_size,ep_size,cp_size", [(8, 1, 1), (4, 2, 1), (1, 1, 8), (2, 1, 4), (2, 2, 2)]
+    )
+    def test_padding_mask_removes_padding_tokens(self, aux_loss_type, tp_size, ep_size, cp_size):
         """Test that padding tokens are correctly excluded from aux loss calculation."""
-        clear_aux_losses_tracker()
+        # Initialize model parallel with given configuration
+        self.setup_model_parallel(tp_size=tp_size, ep_size=ep_size, cp_size=cp_size)
+        
+        try:
+            clear_aux_losses_tracker()
 
-        router = self.new_router(
-            moe_router_load_balancing_type=aux_loss_type,
-            moe_aux_loss_coeff=1.0,
-            moe_router_dtype="fp64",
-        ).cuda()
+            router = self.new_router(
+                moe_router_load_balancing_type=aux_loss_type,
+                moe_aux_loss_coeff=1.0,
+                moe_router_dtype="fp64",
+            ).cuda()
 
-        seq_len = 32
-        batch_size = 2
-        hidden_size = router.config.hidden_size
+            seq_len = 32
+            batch_size = 2
+            hidden_size = router.config.hidden_size
 
-        # Create input with padding
-        hidden_states_full = torch.randn(
-            (seq_len, batch_size, hidden_size), dtype=torch.bfloat16, device='cuda'
-        )
+            # Create input with padding
+            hidden_states_full = torch.randn(
+                (seq_len, batch_size, hidden_size), dtype=torch.bfloat16, device='cuda'
+            )
 
-        # Create padding mask: first half valid, second half padding
-        padding_mask = torch.ones((seq_len, batch_size), dtype=torch.bool, device='cuda')
-        padding_mask[seq_len // 2 :, :] = False
+            # Create padding mask: first half valid, second half padding
+            padding_mask = torch.ones((seq_len, batch_size), dtype=torch.bool, device='cuda')
+            padding_mask[seq_len // 2 :, :] = False
 
-        # Test with padding mask
-        router.weight.grad = None
-        scores_with_mask, routing_map_with_mask = router(
-            hidden_states_full, padding_mask=padding_mask
-        )
-        scores_with_mask.backward(torch.zeros_like(scores_with_mask))
+            # Test with padding mask
+            router.weight.grad = None
+            scores_with_mask, routing_map_with_mask = router(
+                hidden_states_full, padding_mask=padding_mask
+            )
+            scores_with_mask.backward(torch.zeros_like(scores_with_mask))
 
-        loss_name = {
-            "aux_loss": "load_balancing_loss",
-            "seq_aux_loss": "seq_load_balancing_loss",
-            "global_aux_loss": "global_load_balancing_loss",
-        }[aux_loss_type]
+            loss_name = {
+                "aux_loss": "load_balancing_loss",
+                "seq_aux_loss": "seq_load_balancing_loss",
+                "global_aux_loss": "global_load_balancing_loss",
+            }[aux_loss_type]
 
-        tracker = get_moe_layer_wise_logging_tracker()
-        aux_loss_with_mask = tracker[loss_name]["values"][0].clone()
-        grad_with_mask = router.weight.grad.clone()
+            tracker = get_moe_layer_wise_logging_tracker()
+            aux_loss_with_mask = tracker[loss_name]["values"][0].clone()
+            grad_with_mask = router.weight.grad.clone()
 
-        # Test without padding (with only half of the tokens)
-        clear_aux_losses_tracker()
-        router.weight.grad = None
-        hidden_states_valid = hidden_states_full[: seq_len // 2, :, :]
-        scores_without_mask, routing_map_without_mask = router(hidden_states_valid)
-        scores_without_mask.backward(torch.zeros_like(scores_without_mask))
+            # Test without padding (with only half of the tokens)
+            clear_aux_losses_tracker()
+            router.weight.grad = None
+            hidden_states_valid = hidden_states_full[: seq_len // 2, :, :]
+            scores_without_mask, routing_map_without_mask = router(hidden_states_valid)
+            scores_without_mask.backward(torch.zeros_like(scores_without_mask))
 
-        aux_loss_without_mask = tracker[loss_name]["values"][0].clone()
-        grad_without_mask = router.weight.grad.clone()
+            aux_loss_without_mask = tracker[loss_name]["values"][0].clone()
+            grad_without_mask = router.weight.grad.clone()
 
-        # The aux loss with mask should be close to the aux loss without mask
-        torch.testing.assert_close(aux_loss_with_mask, aux_loss_without_mask, rtol=1e-2, atol=1e-3)
-        torch.testing.assert_close(grad_with_mask, grad_without_mask, rtol=1e-2, atol=1e-3)
+            # The aux loss with mask should be close to the aux loss without mask
+            torch.testing.assert_close(aux_loss_with_mask, aux_loss_without_mask, rtol=1e-2, atol=1e-3)
+            torch.testing.assert_close(grad_with_mask, grad_without_mask, rtol=1e-2, atol=1e-3)
 
-        clear_aux_losses_tracker()
+            clear_aux_losses_tracker()
+        finally:
+            # Always cleanup model parallel
+            Utils.destroy_model_parallel()
 
     @pytest.mark.internal
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-    def test_padding_mask_with_z_loss(self):
+    @pytest.mark.parametrize(
+        "tp_size,ep_size,cp_size", [(8, 1, 1), (4, 2, 1), (1, 1, 8), (2, 1, 4), (2, 2, 2)]
+    )
+    def test_padding_mask_with_z_loss(self, tp_size, ep_size, cp_size):
         """Test that padding mask works correctly with z_loss."""
-        clear_aux_losses_tracker()
+        # Initialize model parallel with given configuration
+        self.setup_model_parallel(tp_size=tp_size, ep_size=ep_size, cp_size=cp_size)
+        
+        try:
+            clear_aux_losses_tracker()
 
-        router = self.new_router(
-            moe_router_load_balancing_type="aux_loss",
-            moe_aux_loss_coeff=0.0,
-            moe_z_loss_coeff=1.0,
-            moe_router_dtype="fp32",
-        ).cuda()
+            router = self.new_router(
+                moe_router_load_balancing_type="aux_loss",
+                moe_aux_loss_coeff=0.0,
+                moe_z_loss_coeff=1.0,
+                moe_router_dtype="fp32",
+            ).cuda()
 
-        seq_len = 32
-        batch_size = 2
-        hidden_size = router.config.hidden_size
+            seq_len = 32
+            batch_size = 2
+            hidden_size = router.config.hidden_size
 
-        # Create input
-        hidden_states_full = torch.randn(
-            (seq_len, batch_size, hidden_size), dtype=torch.bfloat16, device='cuda'
-        )
+            # Create input
+            hidden_states_full = torch.randn(
+                (seq_len, batch_size, hidden_size), dtype=torch.bfloat16, device='cuda'
+            )
 
-        # Create padding mask: first half valid, second half padding
-        padding_mask = torch.ones((seq_len, batch_size), dtype=torch.bool, device='cuda')
-        padding_mask[seq_len // 2 :, :] = False
+            # Create padding mask: first half valid, second half padding
+            padding_mask = torch.ones((seq_len, batch_size), dtype=torch.bool, device='cuda')
+            padding_mask[seq_len // 2 :, :] = False
 
-        # Test with padding mask
-        router.weight.grad = None
-        scores_with_mask, _ = router(hidden_states_full, padding_mask=padding_mask)
-        scores_with_mask.sum().backward()
+            # Test with padding mask
+            router.weight.grad = None
+            scores_with_mask, _ = router(hidden_states_full, padding_mask=padding_mask)
+            scores_with_mask.sum().backward()
 
-        tracker = get_moe_layer_wise_logging_tracker()
-        z_loss_with_mask = tracker["z_loss"]["values"][0].clone()
-        grad_with_mask = router.weight.grad.clone()
+            tracker = get_moe_layer_wise_logging_tracker()
+            z_loss_with_mask = tracker["z_loss"]["values"][0].clone()
+            grad_with_mask = router.weight.grad.clone()
 
-        # Test without padding (with only half of the tokens)
-        clear_aux_losses_tracker()
-        router.weight.grad = None
-        hidden_states_valid = hidden_states_full[: seq_len // 2, :, :]
-        scores_without_mask, _ = router(hidden_states_valid)
-        scores_without_mask.sum().backward()
+            # Test without padding (with only half of the tokens)
+            clear_aux_losses_tracker()
+            router.weight.grad = None
+            hidden_states_valid = hidden_states_full[: seq_len // 2, :, :]
+            scores_without_mask, _ = router(hidden_states_valid)
+            scores_without_mask.sum().backward()
 
-        z_loss_without_mask = tracker["z_loss"]["values"][0].clone()
-        grad_without_mask = router.weight.grad.clone()
+            z_loss_without_mask = tracker["z_loss"]["values"][0].clone()
+            grad_without_mask = router.weight.grad.clone()
 
-        # The z_loss with mask should be close to the z_loss without mask
-        torch.testing.assert_close(z_loss_with_mask, z_loss_without_mask, rtol=1e-2, atol=1e-3)
-        torch.testing.assert_close(grad_with_mask, grad_without_mask, rtol=1e-2, atol=1e-3)
+            # The z_loss with mask should be close to the z_loss without mask
+            torch.testing.assert_close(z_loss_with_mask, z_loss_without_mask, rtol=1e-2, atol=1e-3)
+            torch.testing.assert_close(grad_with_mask, grad_without_mask, rtol=1e-2, atol=1e-3)
 
-        clear_aux_losses_tracker()
+            clear_aux_losses_tracker()
+        finally:
+            # Always cleanup model parallel
+            Utils.destroy_model_parallel()

From f1b4e849717fd2be25bddb3b10ab89be9f012874 Mon Sep 17 00:00:00 2001
From: HaochenYuan <haocheny@nvidia.com>
Date: Thu, 13 Nov 2025 01:14:52 -0800
Subject: [PATCH 06/32] fix linting

---
 megatron/core/transformer/moe/router.py           |  3 ++-
 megatron/core/transformer/transformer_block.py    |  7 ++++++-
 tests/unit_tests/transformer/moe/test_aux_loss.py | 10 ++++++----
 3 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py
index 8a897db3ec3..c109f305ae7 100644
--- a/megatron/core/transformer/moe/router.py
+++ b/megatron/core/transformer/moe/router.py
@@ -355,7 +355,8 @@ def _apply_seq_aux_loss(
         # For seq_aux_loss, total_num_tokens should be seq_length (per sequence),
         # not seq_length * batch_size.
         # When padding_mask is not None, compute from tokens_per_expert to get accurate count.
-        # tokens_per_expert shape: [bsz * num_experts], tokens_per_expert.sum() = seq_length * bsz * topk
+        # tokens_per_expert shape: [bsz * num_experts]
+        # tokens_per_expert.sum() = seq_length * bsz * topk
         # Divide by (bsz * topk) to get seq_length
         if padding_mask is not None:
             total_num_tokens = tokens_per_expert.sum() / (bsz * self.topk)
diff --git a/megatron/core/transformer/transformer_block.py b/megatron/core/transformer/transformer_block.py
index 11c4512035d..c31f123979e 100755
--- a/megatron/core/transformer/transformer_block.py
+++ b/megatron/core/transformer/transformer_block.py
@@ -433,7 +433,12 @@ def _checkpointed_forward(
 
         def custom(start: int, end: int):
             def custom_forward(
-                hidden_states, attention_mask, context, context_mask, rotary_pos_emb, padding_mask=None
+                hidden_states,
+                attention_mask,
+                context,
+                context_mask,
+                rotary_pos_emb,
+                padding_mask=None,
             ):
                 for index in range(start, end):
                     layer = self._get_layer(index)
diff --git a/tests/unit_tests/transformer/moe/test_aux_loss.py b/tests/unit_tests/transformer/moe/test_aux_loss.py
index 2cd1da6f381..448a939f1da 100644
--- a/tests/unit_tests/transformer/moe/test_aux_loss.py
+++ b/tests/unit_tests/transformer/moe/test_aux_loss.py
@@ -582,7 +582,7 @@ class TestPaddingMaskAuxLoss:
 
     def setup_model_parallel(self, tp_size=1, ep_size=1, cp_size=1):
         """Initialize model parallel with given configuration.
-        
+
         Args:
             tp_size: Tensor parallel size.
             ep_size: Expert parallel size.
@@ -634,7 +634,7 @@ def test_padding_mask_removes_padding_tokens(self, aux_loss_type, tp_size, ep_si
         """Test that padding tokens are correctly excluded from aux loss calculation."""
         # Initialize model parallel with given configuration
         self.setup_model_parallel(tp_size=tp_size, ep_size=ep_size, cp_size=cp_size)
-        
+
         try:
             clear_aux_losses_tracker()
 
@@ -685,7 +685,9 @@ def test_padding_mask_removes_padding_tokens(self, aux_loss_type, tp_size, ep_si
             grad_without_mask = router.weight.grad.clone()
 
             # The aux loss with mask should be close to the aux loss without mask
-            torch.testing.assert_close(aux_loss_with_mask, aux_loss_without_mask, rtol=1e-2, atol=1e-3)
+            torch.testing.assert_close(
+                aux_loss_with_mask, aux_loss_without_mask, rtol=1e-2, atol=1e-3
+            )
             torch.testing.assert_close(grad_with_mask, grad_without_mask, rtol=1e-2, atol=1e-3)
 
             clear_aux_losses_tracker()
@@ -702,7 +704,7 @@ def test_padding_mask_with_z_loss(self, tp_size, ep_size, cp_size):
         """Test that padding mask works correctly with z_loss."""
         # Initialize model parallel with given configuration
         self.setup_model_parallel(tp_size=tp_size, ep_size=ep_size, cp_size=cp_size)
-        
+
         try:
             clear_aux_losses_tracker()
 

From b97e762bec2e549136fcc5ad555c4e7331f60784 Mon Sep 17 00:00:00 2001
From: HaochenYuan <haocheny@nvidia.com>
Date: Tue, 2 Dec 2025 10:23:22 -0800
Subject: [PATCH 07/32] refactor

---
 megatron/core/transformer/moe/moe_layer.py    | 42 ++++++++--
 megatron/core/transformer/moe/moe_utils.py    | 84 ++++++++++++-------
 megatron/core/transformer/moe/router.py       | 66 +++++++++------
 .../core/transformer/transformer_layer.py     | 75 +++++++----------
 .../transformer/moe/test_aux_loss.py          | 24 +++---
 .../transformer/moe/test_routers.py           |  9 +-
 6 files changed, 173 insertions(+), 127 deletions(-)

diff --git a/megatron/core/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py
index 3a93b9709aa..67e8a0149f9 100644
--- a/megatron/core/transformer/moe/moe_layer.py
+++ b/megatron/core/transformer/moe/moe_layer.py
@@ -184,11 +184,15 @@ def route(self, hidden_states: torch.Tensor, padding_mask: Optional[torch.Tensor
         This method uses the router to determine which experts to send each token to,
         producing routing probabilities and a mapping.
 
+        Note: The padding_mask transformation from [bsz, seq_length] to [seq_length, bsz]
+        is done in the forward() method before calling this method.
+
         Args:
-            hidden_states (torch.Tensor): Input tensor.
-            padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
-                                                   Shape = [seq_length, bsz]. True for valid tokens,
-                                                   False for padding tokens. Defaults to None.
+            hidden_states (torch.Tensor): Input tensor with shape [seq_length, bsz, hidden_size].
+            padding_mask (torch.Tensor, optional): Boolean mask indicating padding positions.
+                Shape = [seq_length, bsz] (after transformation in forward()).
+                True = padding (exclude), False = valid (include).
+                Defaults to None (all tokens are valid).
         """
         probs, routing_map = self.router(hidden_states, padding_mask=padding_mask)
         return probs, routing_map
@@ -285,11 +289,28 @@ def forward(self, hidden_states: torch.Tensor, padding_mask: Optional[torch.Tens
         3. Expert Computation: Experts process the dispatched tokens.
         4. Combine: The outputs from the experts are combined and returned.
 
+        Padding Mask Handling:
+            The padding_mask indicates which positions should be excluded from MoE routing
+            and auxiliary loss computation. This is useful when training with variable-length
+            sequences that are padded to a fixed length.
+
+            Convention: True = padding (exclude), False = valid (include)
+
+            Input Shape Transformation:
+                - Input padding_mask from GPTModel/TransformerLayer: [bsz, seq_length]
+                - Internal MoE processing expects: [seq_length, bsz]
+                - The transformation is done internally in this method
+
+            The mask is passed through to the router, which uses it to:
+            - Exclude padding tokens from auxiliary loss calculations
+            - Exclude padding tokens from expert bias updates
+            - Exclude padding tokens from z-loss computation
+
         Args:
-            hidden_states (torch.Tensor): The input tensor to the MoE layer.
-            padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
-                                                   Shape = [seq_length, bsz]. True for valid tokens,
-                                                   False for padding tokens. Defaults to None.
+            hidden_states (torch.Tensor): The input tensor with shape [seq_length, bsz, hidden_size].
+            padding_mask (torch.Tensor, optional): Boolean mask indicating padding positions.
+                Shape = [bsz, seq_length]. True = padding (exclude), False = valid (include).
+                Defaults to None (all tokens are valid).
 
         Returns:
             A tuple containing the output tensor and the MLP bias, if any.
@@ -300,6 +321,11 @@ def forward(self, hidden_states: torch.Tensor, padding_mask: Optional[torch.Tens
                 "are enabled without also enabling sequence parallelism."
             )
 
+        # Transform padding_mask from [bsz, seq_length] to [seq_length, bsz] if provided
+        # This aligns the mask shape with hidden_states shape [seq_length, bsz, hidden_size]
+        if padding_mask is not None:
+            padding_mask = padding_mask.transpose(0, 1).bool()
+
         # MoE forward: route -> dispatch -> compute -> combine
         def custom_forward(hidden_states, padding_mask=None):
             try:
diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py
index ef435967c8d..92febb9f007 100644
--- a/megatron/core/transformer/moe/moe_utils.py
+++ b/megatron/core/transformer/moe/moe_utils.py
@@ -34,7 +34,7 @@
 _MOE_LAYER_WISE_LOGGING_TRACKER = {}
 
 
-def compute_tokens_per_expert_with_padding(
+def compute_tokens_per_expert(
     routing_map: torch.Tensor,
     padding_mask: Optional[torch.Tensor] = None,
     reshape_for_seq_aux: bool = False,
@@ -42,15 +42,30 @@ def compute_tokens_per_expert_with_padding(
     bsz: Optional[int] = None,
     num_experts: Optional[int] = None,
 ):
-    """Compute tokens_per_expert and total_num_tokens with optional padding mask.
-    This function provides a unified way to compute token counts across different aux loss types.
+    """Compute tokens_per_expert and num_valid_tokens with optional padding mask.
+
+    This function provides a unified way to compute token counts for MoE auxiliary losses.
+    It supports both standard aux_loss/global_aux_loss and sequence-level aux_loss.
+
+    Padding Mask Semantics:
+        The padding_mask follows the standard masking convention where:
+        - True  = padding position (should be EXCLUDED from computation)
+        - False = valid position (should be INCLUDED in computation)
+        This is consistent with common attention mask conventions in transformer models.
+
+    Workflow:
+        1. If padding_mask is provided, create an inverted mask (~padding_mask) to select valid tokens
+        2. Apply the valid mask to routing_map to exclude padding tokens
+        3. Sum the masked routing_map to get tokens_per_expert
+        4. Count valid tokens by summing the inverted mask
 
     Args:
         routing_map (torch.Tensor): Token to expert routing map.
             - For aux_loss/global_aux_loss: shape [num_tokens, num_experts]
             - For seq_aux_loss: shape [num_tokens, num_experts] but will be reshaped
-        padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
-            Shape [num_tokens]. True for valid tokens, False for padding tokens.
+        padding_mask (torch.Tensor, optional): Boolean mask indicating padding positions.
+            Shape [num_tokens]. True = padding (exclude), False = valid (include).
+            Defaults to None (all tokens are valid).
         reshape_for_seq_aux (bool): If True, reshape routing_map for seq_aux_loss computation.
         seq_length (int, optional): Sequence length, required when reshape_for_seq_aux=True.
         bsz (int, optional): Batch size, required when reshape_for_seq_aux=True.
@@ -69,30 +84,34 @@ def compute_tokens_per_expert_with_padding(
         ), "seq_length, bsz, and num_experts must be provided when reshape_for_seq_aux=True"
 
         if padding_mask is not None:
-            # Reshape padding_mask to [seq_length, bsz]
-            padding_mask_reshaped = padding_mask.reshape(seq_length, bsz)
+            # Invert padding_mask: True (padding) -> False, False (valid) -> True
+            valid_mask = ~padding_mask
+            # Reshape valid_mask to [seq_length, bsz]
+            valid_mask_reshaped = valid_mask.reshape(seq_length, bsz)
             # Expand to match routing_map after reshape [seq_length, bsz * num_experts]
             mask_expanded = (
-                padding_mask_reshaped.unsqueeze(-1)
+                valid_mask_reshaped.unsqueeze(-1)
                 .expand(-1, -1, num_experts)
                 .reshape(seq_length, -1)
             )
             routing_map_masked = routing_map.reshape(seq_length, -1) & mask_expanded
             tokens_per_expert = routing_map_masked.sum(dim=0)
-            # Count valid tokens only
-            num_valid_tokens = padding_mask.sum()
+            # Count valid tokens only (where padding_mask is False)
+            num_valid_tokens = valid_mask.sum()
         else:
             tokens_per_expert = routing_map.reshape(seq_length, -1).sum(dim=0)
             num_valid_tokens = routing_map.shape[0]
     else:
         # aux_loss or global_aux_loss
         if padding_mask is not None:
-            # routing_map: [num_tokens, num_experts], padding_mask: [num_tokens]
-            mask_expanded = padding_mask.unsqueeze(-1)
+            # Invert padding_mask: True (padding) -> False, False (valid) -> True
+            valid_mask = ~padding_mask
+            # routing_map: [num_tokens, num_experts], valid_mask: [num_tokens]
+            mask_expanded = valid_mask.unsqueeze(-1)
             routing_map_masked = routing_map & mask_expanded
             tokens_per_expert = routing_map_masked.sum(dim=0)
-            # Count valid tokens only
-            num_valid_tokens = padding_mask.sum()
+            # Count valid tokens only (where padding_mask is False)
+            num_valid_tokens = valid_mask.sum()
         else:
             tokens_per_expert = routing_map.sum(dim=0)
             num_valid_tokens = routing_map.shape[0]
@@ -159,17 +178,17 @@ def switch_load_balancing_loss_func(
         topk (int): The number of experts selected for each token.
         num_experts (int): The number of experts.
         moe_aux_loss_coeff (float): The coefficient for the auxiliary loss.
-        padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
-                                               Shape in [num_tokens]. True for valid tokens,
-                                               False for padding tokens. Defaults to None.
+        padding_mask (torch.Tensor, optional): Boolean mask indicating padding positions.
+                                               Shape [num_tokens]. True = padding (exclude),
+                                               False = valid (include). Defaults to None.
     Returns:
         torch.Tensor: The auxiliary loss for load balancing.
     """
     # Apply padding mask to probs if provided
     if padding_mask is not None:
-        # padding_mask: [num_tokens], probs: [num_tokens, num_experts]
-        mask_expanded = padding_mask.unsqueeze(-1)
-        probs = probs * mask_expanded
+        # Invert padding_mask: True (padding) -> 0, False (valid) -> 1
+        valid_mask = (~padding_mask).unsqueeze(-1)
+        probs = probs * valid_mask
 
     if fused:
         if not HAVE_TE or fused_moe_aux_loss is None:
@@ -197,9 +216,9 @@ def z_loss_func(logits, z_loss_coeff, padding_mask: Optional[torch.Tensor] = Non
     Args:
         logits (torch.Tensor): The logits of the router.
         z_loss_coeff (float): The coefficient for the z-loss.
-        padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
-                                               Shape in [num_tokens]. True for valid tokens,
-                                               False for padding tokens. Defaults to None.
+        padding_mask (torch.Tensor, optional): Boolean mask indicating padding positions.
+                                               Shape [num_tokens]. True = padding (exclude),
+                                               False = valid (include). Defaults to None.
 
     Returns:
         torch.Tensor: The logits after applying the z-loss.
@@ -208,10 +227,12 @@ def z_loss_func(logits, z_loss_coeff, padding_mask: Optional[torch.Tensor] = Non
     z_loss_values = torch.square(logsum)
 
     if padding_mask is not None:
-        # Only compute z_loss for non-padding tokens
-        z_loss_values = z_loss_values * padding_mask
+        # Invert padding_mask: True (padding) -> 0, False (valid) -> 1
+        valid_mask = ~padding_mask
+        # Only compute z_loss for valid (non-padding) tokens
+        z_loss_values = z_loss_values * valid_mask
         # Compute mean over valid tokens only
-        num_valid_tokens = padding_mask.sum()
+        num_valid_tokens = valid_mask.sum()
         z_loss = z_loss_values.sum() / torch.clamp(num_valid_tokens, min=1.0) * z_loss_coeff
     else:
         z_loss = torch.mean(z_loss_values) * z_loss_coeff
@@ -723,9 +744,9 @@ def compute_routing_scores_for_aux_loss(
 
     Args:
         logits (torch.Tensor): The logits tensor after gating, shape: [num_tokens, num_experts].
-        padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
-                                               Shape in [num_tokens]. True for valid tokens,
-                                               False for padding tokens. Defaults to None.
+        padding_mask (torch.Tensor, optional): Boolean mask indicating padding positions.
+                                               Shape [num_tokens]. True = padding (exclude),
+                                               False = valid (include). Defaults to None.
     Returns:
         Tuple[torch.Tensor, torch.Tensor]: routing_map and scores.
     """
@@ -756,8 +777,9 @@ def compute_routing_scores_for_aux_loss(
 
     # Apply padding mask to scores if provided
     if padding_mask is not None:
-        mask_expanded = padding_mask.unsqueeze(-1).to(scores.dtype)
-        scores = scores * mask_expanded
+        # Invert padding_mask: True (padding) -> 0, False (valid) -> 1
+        valid_mask = (~padding_mask).unsqueeze(-1).to(scores.dtype)
+        scores = scores * valid_mask
     return routing_map, scores
 
 
diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py
index c109f305ae7..13065a00dad 100644
--- a/megatron/core/transformer/moe/router.py
+++ b/megatron/core/transformer/moe/router.py
@@ -14,7 +14,7 @@
     apply_random_logits,
     apply_router_token_dropping,
     compute_routing_scores_for_aux_loss,
-    compute_tokens_per_expert_with_padding,
+    compute_tokens_per_expert,
     router_gating_linear,
     save_to_aux_losses_tracker,
     sinkhorn,
@@ -281,17 +281,22 @@ def _apply_aux_loss(
             return probs
 
         # Use unified function to compute tokens_per_expert and num_tokens
-        tokens_per_expert, num_tokens = compute_tokens_per_expert_with_padding(
+        tokens_per_expert, num_tokens = compute_tokens_per_expert(
             routing_map=routing_map, padding_mask=padding_mask, reshape_for_seq_aux=False
         )
 
         tokens_per_expert = reduce_from_tensor_model_parallel_region(
             tokens_per_expert, self.tp_cp_group
         )
-        # When padding_mask is not None, total_num_tokens should be the sum across all ranks.
-        # After all_reduce, tokens_per_expert.sum() gives total token-expert assignments.
-        # Since each token is assigned to topk experts, we divide by topk to get actual token count.
-        # This avoids an extra all_reduce operation for num_tokens.
+        # When padding_mask is not None, we need the global count of valid tokens across all ranks.
+        # Naively, this would require an extra all_reduce on num_valid_tokens.
+        #
+        # Optimization: We derive total_num_tokens from the already all-reduced tokens_per_expert:
+        #   - tokens_per_expert.sum() = total valid token-expert assignments (global)
+        #   - Each valid token is assigned to exactly `topk` experts
+        #   - Therefore: total_valid_tokens = tokens_per_expert.sum() / topk
+        #
+        # This avoids an extra all_reduce operation.
         if padding_mask is not None:
             total_num_tokens = tokens_per_expert.sum() / self.topk
         else:
@@ -340,7 +345,7 @@ def _apply_seq_aux_loss(
         scores_for_aux_loss = scores_for_aux_loss.reshape(seq_length, -1)
 
         # Use unified function to compute tokens_per_expert and num_tokens
-        tokens_per_expert, num_tokens = compute_tokens_per_expert_with_padding(
+        tokens_per_expert, num_tokens = compute_tokens_per_expert(
             routing_map=routing_map,
             padding_mask=padding_mask,
             reshape_for_seq_aux=True,
@@ -376,9 +381,10 @@ def _apply_seq_aux_loss(
             )
             / bsz
         )
-        # Calculate valid token count: seq_length for each batch element
+        # Calculate valid token count: number of non-padding tokens
         if padding_mask is not None:
-            valid_token_count = padding_mask.sum()
+            # Invert padding_mask: count positions where padding_mask is False
+            valid_token_count = (~padding_mask).sum()
         else:
             valid_token_count = seq_length * bsz
 
@@ -405,7 +411,7 @@ def _apply_global_aux_loss(
             return probs
 
         # Use unified function to compute tokens_per_expert and num_tokens
-        tokens_per_expert, num_tokens = compute_tokens_per_expert_with_padding(
+        tokens_per_expert, num_tokens = compute_tokens_per_expert(
             routing_map=routing_map, padding_mask=padding_mask, reshape_for_seq_aux=False
         )
 
@@ -417,10 +423,15 @@ def _apply_global_aux_loss(
         self.ga_steps += 1
         averated_tokens_per_expert = self.global_tokens_per_expert / self.ga_steps
 
-        # When padding_mask is not None, total_num_tokens should be the sum across all ranks.
-        # After all_reduce, tokens_per_expert.sum() gives total token-expert assignments.
-        # Since each token is assigned to topk experts, we divide by topk to get actual token count.
-        # This avoids an extra all_reduce operation for num_tokens.
+        # When padding_mask is not None, we need the global count of valid tokens across all ranks.
+        # Naively, this would require an extra all_reduce on num_valid_tokens.
+        #
+        # Optimization: We derive total_num_tokens from the already all-reduced tokens_per_expert:
+        #   - tokens_per_expert.sum() = total valid token-expert assignments (global)
+        #   - Each valid token is assigned to exactly `topk` experts
+        #   - Therefore: total_valid_tokens = tokens_per_expert.sum() / topk
+        #
+        # This avoids an extra all_reduce operation.
         if padding_mask is not None:
             total_num_tokens = tokens_per_expert.sum() / self.topk
         else:
@@ -502,9 +513,9 @@ def apply_z_loss(self, logits, padding_mask: Optional[torch.Tensor] = None):
 
         Args:
             logits (torch.Tensor): The logits of the router.
-            padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
-                                                   Shape in [num_tokens]. True for valid tokens,
-                                                   False for padding tokens. Defaults to None.
+            padding_mask (torch.Tensor, optional): Boolean mask indicating padding positions.
+                                                   Shape [num_tokens]. True = padding (exclude),
+                                                   False = valid (include). Defaults to None.
 
         Returns:
             torch.Tensor: The logits after applying the z-loss.
@@ -522,7 +533,8 @@ def apply_z_loss(self, logits, padding_mask: Optional[torch.Tensor] = None):
                 # which scales both the main_loss gradient and z_loss gradient by
                 # 1/(num_local_tokens * dp_size * num_micro_batches) in finalize_model_grads().
                 # To correct this scaling, we need to scale the z_loss by num_local_tokens here.
-                num_tokens = padding_mask.sum() if padding_mask is not None else logits.shape[0]
+                # Count valid tokens: sum of inverted mask (False -> True = valid)
+                num_tokens = (~padding_mask).sum() if padding_mask is not None else logits.shape[0]
                 logits = MoEAuxLossAutoScaler.apply(logits, z_loss * num_tokens)
             else:
                 logits = MoEAuxLossAutoScaler.apply(logits, z_loss)
@@ -566,14 +578,14 @@ def _apply_expert_bias(
 
         Args:
             routing_map (torch.Tensor): Token to expert routing map, [num_tokens, num_experts].
-            padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
-                Shape [num_tokens]. True for valid tokens, False for padding tokens.
+            padding_mask (torch.Tensor, optional): Boolean mask indicating padding positions.
+                Shape [num_tokens]. True = padding (exclude), False = valid (include).
         """
         if self.enable_expert_bias and torch.is_grad_enabled():
             with torch.no_grad():
                 if padding_mask is not None:
                     # Only count valid tokens, exclude padding tokens
-                    # This logic is consistent with compute_tokens_per_expert_with_padding
+                    # This logic is consistent with compute_tokens_per_expert
                     mask_expanded = padding_mask.unsqueeze(-1)
                     routing_map_masked = routing_map & mask_expanded
                     self.local_tokens_per_expert += routing_map_masked.sum(dim=0)
@@ -585,9 +597,9 @@ def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = N
 
         Args:
             logits (torch.Tensor): Logits tensor after gating.
-            padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
-                                                   Shape = [seq_length, bsz]. True for valid tokens,
-                                                   False for padding tokens. Defaults to None.
+            padding_mask (torch.Tensor, optional): Boolean mask indicating padding positions.
+                                                   Shape = [seq_length, bsz]. True = padding (exclude),
+                                                   False = valid (include). Defaults to None.
 
         Returns:
             probs (torch.Tensor): The probabilities of token to experts assignment.
@@ -675,9 +687,9 @@ def forward(self, input: torch.Tensor, padding_mask: Optional[torch.Tensor] = No
 
         Args:
             input (torch.Tensor): Input tensor.
-            padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
-                                                   Shape = [seq_length, bsz]. True for valid tokens,
-                                                   False for padding tokens. Defaults to None.
+            padding_mask (torch.Tensor, optional): Boolean mask indicating padding positions.
+                                                   Shape = [seq_length, bsz]. True = padding (exclude),
+                                                   False = valid (include). Defaults to None.
         """
         self._maintain_float32_expert_bias()
 
diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py
index 49eb7f2b9ec..f659b54543a 100644
--- a/megatron/core/transformer/transformer_layer.py
+++ b/megatron/core/transformer/transformer_layer.py
@@ -600,7 +600,9 @@ def _forward_mlp(self, hidden_states, inference_context=None, padding_mask=None)
         Args:
             hidden_states (Tensor): Transformed hidden states before the MLP layernorm.
             inference_context: Inference context for optimizations.
-            padding_mask (Tensor, optional): Padding mask for MoE routing. Shape [bsz, seq_length].
+            padding_mask (Tensor, optional): Padding mask for MoE routing.
+                Shape [bsz, seq_length]. True = padding (exclude), False = valid (include).
+                Only used for MoE layers to exclude padding tokens from routing computations.
 
         Returns:
             output (Tensor): Transformed hidden states of shape [s, b, h].
@@ -614,12 +616,6 @@ def _forward_mlp(self, hidden_states, inference_context=None, padding_mask=None)
         # Residual connection.
         residual = hidden_states
 
-        padding_mask_for_moe = None
-        if self.is_moe_layer and padding_mask is not None:
-            # padding_mask from GPTModel: [bsz, seq_length]
-            # MoE layer needs: [seq_length, bsz]
-            padding_mask_for_moe = padding_mask.transpose(0, 1).bool()
-
         if self.offload_mlp_norm:
             hidden_states = fine_grained_offloading_group_start(hidden_states, name="mlp_norm")
         # Optional Layer norm post the cross-attention.
@@ -642,6 +638,20 @@ def _forward_mlp(self, hidden_states, inference_context=None, padding_mask=None)
             and not isinstance(self.mlp, IdentityOp)
         )
 
+        # Helper function for MoE forward with padding_mask
+        def get_mlp_forward_fn():
+            """Return the appropriate MLP forward function based on layer type."""
+            if self.is_moe_layer and padding_mask is not None:
+                def mlp_forward(hidden_states):
+                    return self.mlp(hidden_states, padding_mask=padding_mask)
+                return mlp_forward
+            elif self.is_moe_layer:
+                def mlp_forward(hidden_states):
+                    return self.mlp(hidden_states)
+                return mlp_forward
+            else:
+                return self.mlp
+
         if (
             self.is_moe_layer
             and self.config.cuda_graph_impl == "transformer_engine"
@@ -652,48 +662,25 @@ def _forward_mlp(self, hidden_states, inference_context=None, padding_mask=None)
             assert (
                 not self.recompute_pre_mlp_layernorm
             ), "Recomputation is not supported for CUDA graph."
-            cudagraph_outputs = self.mlp(
-                pre_mlp_layernorm_output, padding_mask=padding_mask_for_moe
-            )
+            cudagraph_outputs = self.mlp(pre_mlp_layernorm_output, padding_mask=padding_mask)
             return cudagraph_outputs + [residual]
         elif self.recompute_mlp:
+            mlp_fn = get_mlp_forward_fn()
             if self.config.fp8:
                 # import here to avoid circular import
                 from megatron.core.extensions.transformer_engine import te_checkpoint
 
-                if self.is_moe_layer:
-                    # For MoE with checkpointing, we need a wrapper to pass padding_mask
-                    def mlp_forward_with_padding(hidden_states):
-                        return self.mlp(hidden_states, padding_mask=padding_mask_for_moe)
-
-                    mlp_output_with_bias = te_checkpoint(
-                        mlp_forward_with_padding,
-                        False,
-                        tensor_parallel.random.get_cuda_rng_tracker,
-                        self.pg_collection.tp,
-                        pre_mlp_layernorm_output,
-                    )
-                else:
-                    mlp_output_with_bias = te_checkpoint(
-                        self.mlp,
-                        False,
-                        tensor_parallel.random.get_cuda_rng_tracker,
-                        self.pg_collection.tp,
-                        pre_mlp_layernorm_output,
-                    )
+                mlp_output_with_bias = te_checkpoint(
+                    mlp_fn,
+                    False,
+                    tensor_parallel.random.get_cuda_rng_tracker,
+                    self.pg_collection.tp,
+                    pre_mlp_layernorm_output,
+                )
             else:
-                if self.is_moe_layer:
-                    # For MoE with checkpointing, we need a wrapper to pass padding_mask
-                    def mlp_forward_with_padding(hidden_states):
-                        return self.mlp(hidden_states, padding_mask=padding_mask_for_moe)
-
-                    mlp_output_with_bias = tensor_parallel.checkpoint(
-                        mlp_forward_with_padding, False, pre_mlp_layernorm_output
-                    )
-                else:
-                    mlp_output_with_bias = tensor_parallel.checkpoint(
-                        self.mlp, False, pre_mlp_layernorm_output
-                    )
+                mlp_output_with_bias = tensor_parallel.checkpoint(
+                    mlp_fn, False, pre_mlp_layernorm_output
+                )
         elif should_chunk_mlp_for_prefill:
             # Chunk input along sequence dimension
             num_chunks = min(self.config.mlp_chunks_for_prefill, pre_mlp_layernorm_output.shape[0])
@@ -709,9 +696,7 @@ def mlp_forward_with_padding(hidden_states):
             mlp_output_with_bias = (mlp_output, bias_output)
         else:
             if self.is_moe_layer:
-                mlp_output_with_bias = self.mlp(
-                    pre_mlp_layernorm_output, padding_mask=padding_mask_for_moe
-                )
+                mlp_output_with_bias = self.mlp(pre_mlp_layernorm_output, padding_mask=padding_mask)
             else:
                 mlp_output_with_bias = self.mlp(pre_mlp_layernorm_output)
 
diff --git a/tests/unit_tests/transformer/moe/test_aux_loss.py b/tests/unit_tests/transformer/moe/test_aux_loss.py
index 448a939f1da..db40269579a 100644
--- a/tests/unit_tests/transformer/moe/test_aux_loss.py
+++ b/tests/unit_tests/transformer/moe/test_aux_loss.py
@@ -653,9 +653,10 @@ def test_padding_mask_removes_padding_tokens(self, aux_loss_type, tp_size, ep_si
                 (seq_len, batch_size, hidden_size), dtype=torch.bfloat16, device='cuda'
             )
 
-            # Create padding mask: first half valid, second half padding
-            padding_mask = torch.ones((seq_len, batch_size), dtype=torch.bool, device='cuda')
-            padding_mask[seq_len // 2 :, :] = False
+            # Create padding mask: first half valid (False), second half padding (True)
+            # Convention: True = padding (exclude), False = valid (include)
+            padding_mask = torch.zeros((seq_len, batch_size), dtype=torch.bool, device='cuda')
+            padding_mask[seq_len // 2 :, :] = True
 
             # Test with padding mask
             router.weight.grad = None
@@ -685,10 +686,8 @@ def test_padding_mask_removes_padding_tokens(self, aux_loss_type, tp_size, ep_si
             grad_without_mask = router.weight.grad.clone()
 
             # The aux loss with mask should be close to the aux loss without mask
-            torch.testing.assert_close(
-                aux_loss_with_mask, aux_loss_without_mask, rtol=1e-2, atol=1e-3
-            )
-            torch.testing.assert_close(grad_with_mask, grad_without_mask, rtol=1e-2, atol=1e-3)
+            assert torch.equal(aux_loss_with_mask, aux_loss_without_mask)
+            assert torch.equal(grad_with_mask, grad_without_mask)
 
             clear_aux_losses_tracker()
         finally:
@@ -724,9 +723,10 @@ def test_padding_mask_with_z_loss(self, tp_size, ep_size, cp_size):
                 (seq_len, batch_size, hidden_size), dtype=torch.bfloat16, device='cuda'
             )
 
-            # Create padding mask: first half valid, second half padding
-            padding_mask = torch.ones((seq_len, batch_size), dtype=torch.bool, device='cuda')
-            padding_mask[seq_len // 2 :, :] = False
+            # Create padding mask: first half valid (False), second half padding (True)
+            # Convention: True = padding (exclude), False = valid (include)
+            padding_mask = torch.zeros((seq_len, batch_size), dtype=torch.bool, device='cuda')
+            padding_mask[seq_len // 2 :, :] = True
 
             # Test with padding mask
             router.weight.grad = None
@@ -748,8 +748,8 @@ def test_padding_mask_with_z_loss(self, tp_size, ep_size, cp_size):
             grad_without_mask = router.weight.grad.clone()
 
             # The z_loss with mask should be close to the z_loss without mask
-            torch.testing.assert_close(z_loss_with_mask, z_loss_without_mask, rtol=1e-2, atol=1e-3)
-            torch.testing.assert_close(grad_with_mask, grad_without_mask, rtol=1e-2, atol=1e-3)
+            assert torch.equal(z_loss_with_mask, z_loss_without_mask)
+            assert torch.equal(grad_with_mask, grad_without_mask)
 
             clear_aux_losses_tracker()
         finally:
diff --git a/tests/unit_tests/transformer/moe/test_routers.py b/tests/unit_tests/transformer/moe/test_routers.py
index d11aef207c4..abd1a4db2dc 100644
--- a/tests/unit_tests/transformer/moe/test_routers.py
+++ b/tests/unit_tests/transformer/moe/test_routers.py
@@ -137,10 +137,11 @@ def test_router_with_padding_mask(self):
         # Create input with shape [seq_len, batch_size, hidden_size]
         hidden_states = torch.randn((seq_len, batch_size, hidden_size)).cuda().bfloat16()
 
-        # Create padding mask: first half valid, second half padding
+        # Create padding mask: first half valid (False), second half padding (True)
         # padding_mask shape: [seq_len, batch_size]
-        padding_mask = torch.ones((seq_len, batch_size), dtype=torch.bool, device='cuda')
-        padding_mask[seq_len // 2 :, :] = False  # Second half is padding
+        # Convention: True = padding (exclude), False = valid (include)
+        padding_mask = torch.zeros((seq_len, batch_size), dtype=torch.bool, device='cuda')
+        padding_mask[seq_len // 2 :, :] = True  # Second half is padding
 
         # Test forward pass with padding mask
         with torch.no_grad():
@@ -169,7 +170,7 @@ def test_router_with_padding_mask(self):
             )
 
             # Verify that probs for valid tokens are similar
-            torch.testing.assert_close(probs_valid_part, probs_without_mask, rtol=1e-3, atol=1e-3)
+            assert torch.equal(probs_valid_part, probs_without_mask)
 
     @pytest.mark.internal
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")

From 80a06dea00d209385afc401893b8bc8bc9062a61 Mon Sep 17 00:00:00 2001
From: HaochenYuan <haocheny@nvidia.com>
Date: Tue, 2 Dec 2025 10:26:14 -0800
Subject: [PATCH 08/32] refactor

---
 megatron/core/transformer/moe/moe_layer.py     | 2 +-
 megatron/core/transformer/moe/moe_utils.py     | 2 +-
 megatron/core/transformer/moe/router.py        | 4 ++--
 megatron/core/transformer/transformer_layer.py | 4 ++++
 4 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/megatron/core/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py
index 67e8a0149f9..8cc85b25c1e 100644
--- a/megatron/core/transformer/moe/moe_layer.py
+++ b/megatron/core/transformer/moe/moe_layer.py
@@ -307,7 +307,7 @@ def forward(self, hidden_states: torch.Tensor, padding_mask: Optional[torch.Tens
             - Exclude padding tokens from z-loss computation
 
         Args:
-            hidden_states (torch.Tensor): The input tensor with shape [seq_length, bsz, hidden_size].
+            hidden_states (torch.Tensor): The input tensor shape is [seq_length, bsz, hidden_size].
             padding_mask (torch.Tensor, optional): Boolean mask indicating padding positions.
                 Shape = [bsz, seq_length]. True = padding (exclude), False = valid (include).
                 Defaults to None (all tokens are valid).
diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py
index 92febb9f007..12cb7e7477a 100644
--- a/megatron/core/transformer/moe/moe_utils.py
+++ b/megatron/core/transformer/moe/moe_utils.py
@@ -54,7 +54,7 @@ def compute_tokens_per_expert(
         This is consistent with common attention mask conventions in transformer models.
 
     Workflow:
-        1. If padding_mask is provided, create an inverted mask (~padding_mask) to select valid tokens
+        1. If padding_mask is provided, create an inverted mask (~padding_mask) for valid tokens
         2. Apply the valid mask to routing_map to exclude padding tokens
         3. Sum the masked routing_map to get tokens_per_expert
         4. Count valid tokens by summing the inverted mask
diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py
index 13065a00dad..379ff4c9dd2 100644
--- a/megatron/core/transformer/moe/router.py
+++ b/megatron/core/transformer/moe/router.py
@@ -598,7 +598,7 @@ def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = N
         Args:
             logits (torch.Tensor): Logits tensor after gating.
             padding_mask (torch.Tensor, optional): Boolean mask indicating padding positions.
-                                                   Shape = [seq_length, bsz]. True = padding (exclude),
+                                                   Shape = [seq_length, bsz]. True = padding,
                                                    False = valid (include). Defaults to None.
 
         Returns:
@@ -688,7 +688,7 @@ def forward(self, input: torch.Tensor, padding_mask: Optional[torch.Tensor] = No
         Args:
             input (torch.Tensor): Input tensor.
             padding_mask (torch.Tensor, optional): Boolean mask indicating padding positions.
-                                                   Shape = [seq_length, bsz]. True = padding (exclude),
+                                                   Shape = [seq_length, bsz]. True = padding,
                                                    False = valid (include). Defaults to None.
         """
         self._maintain_float32_expert_bias()
diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py
index f659b54543a..5ebb3753422 100644
--- a/megatron/core/transformer/transformer_layer.py
+++ b/megatron/core/transformer/transformer_layer.py
@@ -642,12 +642,16 @@ def _forward_mlp(self, hidden_states, inference_context=None, padding_mask=None)
         def get_mlp_forward_fn():
             """Return the appropriate MLP forward function based on layer type."""
             if self.is_moe_layer and padding_mask is not None:
+
                 def mlp_forward(hidden_states):
                     return self.mlp(hidden_states, padding_mask=padding_mask)
+
                 return mlp_forward
             elif self.is_moe_layer:
+
                 def mlp_forward(hidden_states):
                     return self.mlp(hidden_states)
+
                 return mlp_forward
             else:
                 return self.mlp

From 7664c4d7b9b55369034c00f19ae2c5a416dd4834 Mon Sep 17 00:00:00 2001
From: HaochenYuan <haocheny@nvidia.com>
Date: Tue, 2 Dec 2025 10:35:59 -0800
Subject: [PATCH 09/32] add comments for padding mask

---
 megatron/core/models/gpt/gpt_model.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py
index 3f9f34f9b59..452194e91dd 100644
--- a/megatron/core/models/gpt/gpt_model.py
+++ b/megatron/core/models/gpt/gpt_model.py
@@ -456,6 +456,9 @@ def forward(
         Args:
             runtime_gather_output (bool): Gather output at runtime. Default None means
                 `parallel_output` arg in the constructor will be used.
+            padding_mask (Tensor, optional): Padding mask for MoE routing.
+                Shape [bsz, seq_length]. True = padding (exclude), False = valid (include).
+                Only used for MoE layers to exclude padding tokens from routing computations.
         """
         if self.config.fine_grained_activation_offloading:
             self.preprocess_for_fine_grained_offloading()

From 8acb7e90d1ba62fc91189d2c1aabda7409d4aacd Mon Sep 17 00:00:00 2001
From: HaochenYuan <haocheny@nvidia.com>
Date: Tue, 2 Dec 2025 23:54:09 -0800
Subject: [PATCH 10/32] refactor: simplify logic

---
 megatron/core/transformer/moe/moe_layer.py    |  42 +-
 megatron/core/transformer/moe/moe_utils.py    | 461 +++++++++++++++---
 megatron/core/transformer/moe/router.py       |  36 +-
 .../core/transformer/transformer_layer.py     |  17 +-
 4 files changed, 466 insertions(+), 90 deletions(-)

diff --git a/megatron/core/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py
index 8cc85b25c1e..bca90c7f916 100644
--- a/megatron/core/transformer/moe/moe_layer.py
+++ b/megatron/core/transformer/moe/moe_layer.py
@@ -12,6 +12,7 @@
 from megatron.core.transformer.moe.moe_utils import (
     MoECudaGraphPartialCaptureSignal,
     MoECudaGraphTensorStore,
+    PaddingMaskInfo,
     get_default_pg_collection,
     maybe_skip_or_early_return_by_cudagraph,
 )
@@ -290,24 +291,30 @@ def forward(self, hidden_states: torch.Tensor, padding_mask: Optional[torch.Tens
         4. Combine: The outputs from the experts are combined and returned.
 
         Padding Mask Handling:
-            The padding_mask indicates which positions should be excluded from MoE routing
-            and auxiliary loss computation. This is useful when training with variable-length
-            sequences that are padded to a fixed length.
+        =====================
+        The padding_mask indicates which positions should be excluded from MoE routing
+        and auxiliary loss computation. This is useful when training with variable-length
+        sequences that are padded to a fixed length.
 
-            Convention: True = padding (exclude), False = valid (include)
+        Padding Mask Semantics:
+            - True  = padding position (should be EXCLUDED from computation)
+            - False = valid position (should be INCLUDED in computation)
 
-            Input Shape Transformation:
-                - Input padding_mask from GPTModel/TransformerLayer: [bsz, seq_length]
-                - Internal MoE processing expects: [seq_length, bsz]
-                - The transformation is done internally in this method
+        Shape Transformation:
+            - Input padding_mask from GPTModel/TransformerLayer: [bsz, seq_length]
+            - This method uses PaddingMaskInfo to transform to: [seq_length, bsz]
+            - This aligns with hidden_states shape [seq_length, bsz, hidden_size]
 
-            The mask is passed through to the router, which uses it to:
+        The mask is passed through to the router, which uses utility functions to:
             - Exclude padding tokens from auxiliary loss calculations
             - Exclude padding tokens from expert bias updates
             - Exclude padding tokens from z-loss computation
 
+        Note: The mask does NOT affect actual routing decisions - padding tokens are
+        still routed to experts, but they don't contribute to loss calculations.
+
         Args:
-            hidden_states (torch.Tensor): The input tensor shape is [seq_length, bsz, hidden_size].
+            hidden_states (torch.Tensor): The input tensor shape [seq_length, bsz, hidden_size].
             padding_mask (torch.Tensor, optional): Boolean mask indicating padding positions.
                 Shape = [bsz, seq_length]. True = padding (exclude), False = valid (include).
                 Defaults to None (all tokens are valid).
@@ -321,10 +328,11 @@ def forward(self, hidden_states: torch.Tensor, padding_mask: Optional[torch.Tens
                 "are enabled without also enabling sequence parallelism."
             )
 
-        # Transform padding_mask from [bsz, seq_length] to [seq_length, bsz] if provided
-        # This aligns the mask shape with hidden_states shape [seq_length, bsz, hidden_size]
-        if padding_mask is not None:
-            padding_mask = padding_mask.transpose(0, 1).bool()
+        # Use PaddingMaskInfo for unified mask handling
+        # Transform from [bsz, seq_length] to [seq_length, bsz] format
+        seq_length, bsz = hidden_states.shape[:2]
+        mask_info = PaddingMaskInfo.from_bsz_seq(padding_mask, seq_length, bsz)
+        padding_mask_seq_bsz = mask_info.seq_bsz
 
         # MoE forward: route -> dispatch -> compute -> combine
         def custom_forward(hidden_states, padding_mask=None):
@@ -353,14 +361,14 @@ def custom_forward(hidden_states, padding_mask=None):
                     tensor_parallel.random.get_cuda_rng_tracker,
                     parallel_state.get_tensor_model_parallel_group(),
                     hidden_states,
-                    padding_mask,
+                    padding_mask_seq_bsz,
                 )
             else:
                 outputs = tensor_parallel.checkpoint(
-                    custom_forward, False, hidden_states, padding_mask
+                    custom_forward, False, hidden_states, padding_mask_seq_bsz
                 )
         else:
-            outputs = custom_forward(hidden_states, padding_mask)
+            outputs = custom_forward(hidden_states, padding_mask_seq_bsz)
 
         return outputs
 
diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py
index 12cb7e7477a..aa3b15d9c79 100644
--- a/megatron/core/transformer/moe/moe_utils.py
+++ b/megatron/core/transformer/moe/moe_utils.py
@@ -1,5 +1,46 @@
 # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
+"""
+MoE Utility Functions
+=====================
+
+This module provides utilities for Mixture of Experts (MoE) routing and loss computation.
+
+Padding Mask Handling
+---------------------
+MoE models need special handling for padding tokens in variable-length sequences.
+Without proper masking, padding tokens would:
+1. Artificially inflate expert load statistics
+2. Contribute to z-loss despite being meaningless
+3. Skew load balancing metrics
+
+This module provides unified utilities for padding mask handling:
+
+Padding Mask Semantics (IMPORTANT):
+    The padding_mask follows the standard masking convention where:
+    - True  = padding position (should be EXCLUDED from computation)
+    - False = valid position (should be INCLUDED in computation)
+    This is consistent with common attention mask conventions in transformer models.
+
+Shape Conventions:
+    - User input (GPTModel): [batch_size, seq_length]
+    - Internal format: [seq_length, batch_size] (Megatron convention)
+    - Flattened: [num_tokens] where num_tokens = seq_length * batch_size
+
+Key Classes and Functions:
+    - PaddingMaskInfo: Unified container for padding mask with different shape views
+    - apply_padding_mask_to_routing_map(): Apply mask to routing decisions
+    - apply_padding_mask_to_scores(): Zero out scores for padding tokens
+    - compute_tokens_per_expert(): Unified token counting with mask support
+
+Example Workflow:
+    1. User provides padding_mask [bsz, seq] to GPTModel.forward()
+    2. MoELayer transposes to [seq, bsz] to match hidden_states
+    3. Router creates PaddingMaskInfo for unified handling
+    4. Router flattens mask to [num_tokens] for per-token operations
+    5. Aux loss functions use mask utilities to exclude padding from statistics
+"""
+
 import math
 from dataclasses import dataclass
 from typing import List, Optional, Union
@@ -30,6 +71,295 @@
     HAVE_TE = False
 
 
+from dataclasses import dataclass
+from typing import Tuple
+
+
+@dataclass
+class PaddingMaskInfo:
+    """Container for padding mask with different views/shapes.
+
+    This class provides a unified interface for handling padding masks throughout
+    the MoE routing and auxiliary loss computation pipeline.
+
+    Padding Mask Semantics:
+        - True  = padding position (should be EXCLUDED from computation)
+        - False = valid position (should be INCLUDED in computation)
+
+    The padding mask workflow in MoE:
+    ================================
+    1. User provides padding_mask to GPTModel.forward() with shape [batch_size, seq_length]
+       - True = padding token (exclude), False = valid token (include)
+
+    2. MoELayer transposes it to [seq_length, batch_size] to match hidden_states layout
+       (Megatron uses [seq, batch, hidden] format internally)
+
+    3. Router flattens it to [num_tokens] where num_tokens = seq_length * batch_size
+
+    4. For operations with routing_map [num_tokens, num_experts], we use utility
+       functions that handle mask expansion internally
+
+    Shape Conventions:
+    -----------------
+    - bsz_seq: [batch_size, seq_length] - User input format
+    - seq_bsz: [seq_length, batch_size] - Internal Megatron format
+    - flat: [num_tokens] - Flattened for per-token operations
+
+    Attributes:
+        seq_bsz: Mask in [seq_length, batch_size] format (None if no mask provided)
+        flat: Flattened mask in [num_tokens] format
+        num_valid_tokens: Count of valid (non-padding) tokens
+        seq_length: Sequence length
+        batch_size: Batch size
+
+    Example Usage:
+    -------------
+    >>> # Create from user-provided mask [bsz, seq]
+    >>> mask_info = PaddingMaskInfo.from_bsz_seq(padding_mask, seq_length=32, batch_size=4)
+    >>> # Check if mask exists
+    >>> if mask_info.has_mask:
+    ...     # Use flattened mask in aux loss computation
+    ...     tokens_per_expert, num_valid = apply_padding_mask_to_routing_map(
+    ...         routing_map, mask_info.flat
+    ...     )
+    """
+
+    # Mask in [seq_length, batch_size] format (None if no mask provided)
+    seq_bsz: Optional[torch.Tensor]
+
+    # Flattened mask in [num_tokens] format
+    flat: Optional[torch.Tensor]
+
+    # Number of valid (non-padding) tokens
+    num_valid_tokens: Optional[torch.Tensor]
+
+    # Metadata
+    seq_length: int
+    batch_size: int
+
+    @classmethod
+    def from_bsz_seq(
+        cls, padding_mask: Optional[torch.Tensor], seq_length: int, batch_size: int
+    ) -> "PaddingMaskInfo":
+        """Create PaddingMaskInfo from user-provided mask in [batch_size, seq_length] format.
+
+        Args:
+            padding_mask: Boolean tensor [batch_size, seq_length].
+                         True = padding (exclude), False = valid (include).
+                         If None, creates an "all valid" mask info.
+            seq_length: Sequence length.
+            batch_size: Batch size.
+
+        Returns:
+            PaddingMaskInfo with all shape variants populated.
+        """
+        if padding_mask is None:
+            return cls(
+                seq_bsz=None,
+                flat=None,
+                num_valid_tokens=None,
+                seq_length=seq_length,
+                batch_size=batch_size,
+            )
+
+        # Transpose to [seq_length, batch_size] for internal use
+        seq_bsz = padding_mask.transpose(0, 1).bool()
+
+        # Flatten to [num_tokens]
+        flat = seq_bsz.reshape(-1)
+
+        # Count valid tokens (where padding_mask is False)
+        num_valid_tokens = (~flat).sum()
+
+        return cls(
+            seq_bsz=seq_bsz,
+            flat=flat,
+            num_valid_tokens=num_valid_tokens,
+            seq_length=seq_length,
+            batch_size=batch_size,
+        )
+
+    @classmethod
+    def from_seq_bsz(
+        cls, padding_mask: Optional[torch.Tensor], seq_length: int, batch_size: int
+    ) -> "PaddingMaskInfo":
+        """Create PaddingMaskInfo from mask already in [seq_length, batch_size] format.
+
+        Use this when the mask has already been transposed (e.g., inside MoELayer).
+
+        Args:
+            padding_mask: Boolean tensor [seq_length, batch_size].
+                         True = padding (exclude), False = valid (include).
+                         If None, creates an "all valid" mask info.
+            seq_length: Sequence length.
+            batch_size: Batch size.
+
+        Returns:
+            PaddingMaskInfo with all shape variants populated.
+        """
+        if padding_mask is None:
+            return cls(
+                seq_bsz=None,
+                flat=None,
+                num_valid_tokens=None,
+                seq_length=seq_length,
+                batch_size=batch_size,
+            )
+
+        seq_bsz = padding_mask.bool()
+        flat = seq_bsz.reshape(-1)
+        # Count valid tokens (where padding_mask is False)
+        num_valid_tokens = (~flat).sum()
+
+        return cls(
+            seq_bsz=seq_bsz,
+            flat=flat,
+            num_valid_tokens=num_valid_tokens,
+            seq_length=seq_length,
+            batch_size=batch_size,
+        )
+
+    @property
+    def has_mask(self) -> bool:
+        """Check if a padding mask was provided."""
+        return self.flat is not None
+
+    @property
+    def num_total_tokens(self) -> int:
+        """Total number of tokens (including padding)."""
+        return self.seq_length * self.batch_size
+
+
+def apply_padding_mask_to_routing_map(
+    routing_map: torch.Tensor, padding_mask_flat: Optional[torch.Tensor]
+) -> Tuple[torch.Tensor, Union[int, torch.Tensor]]:
+    """Apply padding mask to routing map and compute tokens per expert.
+
+    This is a unified function to handle the common pattern of:
+    1. Inverting padding_mask (True=padding -> False, False=valid -> True)
+    2. Expanding the valid mask from [num_tokens] to [num_tokens, 1]
+    3. Masking routing_map with the expanded mask
+    4. Computing tokens_per_expert from the masked routing_map
+
+    Args:
+        routing_map: Boolean tensor [num_tokens, num_experts] indicating token-expert assignments.
+        padding_mask_flat: Boolean tensor [num_tokens].
+            True = padding (exclude), False = valid (include).
+            If None, no masking is applied.
+
+    Returns:
+        Tuple of:
+        - tokens_per_expert: Tensor [num_experts] with count of valid tokens per expert
+        - num_valid_tokens: Scalar tensor or int with total valid token count
+
+    Example:
+        >>> routing_map = torch.tensor([[True, False], [False, True], [True, False]])
+        >>> # Last token is padding (True = exclude)
+        >>> padding_mask = torch.tensor([False, False, True])
+        >>> tokens_per_expert, num_valid
+            = apply_padding_mask_to_routing_map(routing_map, padding_mask)
+        >>> print(tokens_per_expert)  # tensor([1, 1]) - only counts first two tokens
+        >>> print(num_valid)  # tensor(2)
+    """
+    if padding_mask_flat is None:
+        tokens_per_expert = routing_map.sum(dim=0)
+        num_valid_tokens = routing_map.shape[0]
+    else:
+        # Invert: True (padding) -> False, False (valid) -> True
+        valid_mask = ~padding_mask_flat
+        # Expand mask for broadcasting: [num_tokens] -> [num_tokens, 1]
+        mask_expanded = valid_mask.unsqueeze(-1)
+        # Apply mask: only count tokens where valid_mask is True
+        routing_map_masked = routing_map & mask_expanded
+        tokens_per_expert = routing_map_masked.sum(dim=0)
+        num_valid_tokens = valid_mask.sum()
+
+    return tokens_per_expert, num_valid_tokens
+
+
+def apply_padding_mask_to_scores(
+    scores: torch.Tensor, padding_mask_flat: Optional[torch.Tensor]
+) -> torch.Tensor:
+    """Zero out scores for padding tokens.
+
+    Args:
+        scores: Tensor [num_tokens, num_experts] with routing scores/probabilities.
+        padding_mask_flat: Boolean tensor [num_tokens].
+            True = padding (exclude), False = valid (include).
+            If None, no masking is applied.
+
+    Returns:
+        Masked scores with padding positions set to 0.
+
+    Example:
+        >>> scores = torch.tensor([[0.5, 0.5], [0.3, 0.7], [0.8, 0.2]])
+        >>> # Last token is padding (True = exclude)
+        >>> padding_mask = torch.tensor([False, False, True])
+        >>> masked_scores = apply_padding_mask_to_scores(scores, padding_mask)
+        >>> print(masked_scores)  # Last row is zeroed out
+    """
+    if padding_mask_flat is None:
+        return scores
+
+    # Invert: True (padding) -> 0, False (valid) -> 1
+    valid_mask = (~padding_mask_flat).unsqueeze(-1).to(scores.dtype)
+    return scores * valid_mask
+
+
+def _compute_tokens_for_seq_aux_loss(
+    routing_map: torch.Tensor,
+    padding_mask: Optional[torch.Tensor],
+    seq_length: int,
+    bsz: int,
+    num_experts: int,
+) -> Tuple[torch.Tensor, Union[int, torch.Tensor]]:
+    """Internal helper for seq_aux_loss token counting.
+
+    For seq_aux_loss, we need per-sequence, per-expert token counts.
+    This requires reshaping the data to separate batch elements.
+
+    Reshape logic:
+    - routing_map [num_tokens, num_experts] -> [seq_length, bsz, num_experts]
+                                            -> [seq_length, bsz * num_experts]
+    - padding_mask [num_tokens] -> [seq_length, bsz] -> broadcast to [seq_length, bsz * num_experts]
+
+    Args:
+        routing_map: Token to expert routing map [num_tokens, num_experts].
+        padding_mask: Boolean mask [num_tokens]. True = padding (exclude), False = valid (include).
+        seq_length: Sequence length.
+        bsz: Batch size.
+        num_experts: Number of experts.
+
+    Returns:
+        Tuple of (tokens_per_expert, num_valid_tokens):
+        - tokens_per_expert: [bsz * num_experts]
+        - num_valid_tokens: Scalar count of valid tokens
+    """
+    if padding_mask is not None:
+        # Invert: True (padding) -> False, False (valid) -> True
+        valid_mask = ~padding_mask
+        # Reshape valid_mask: [num_tokens] -> [seq_length, bsz]
+        valid_mask_2d = valid_mask.reshape(seq_length, bsz)
+
+        # Expand for broadcasting with routing_map
+        # [seq_length, bsz] -> [seq_length, bsz, 1] -> [seq_length, bsz, num_experts]
+        # -> [seq_length, bsz * num_experts]
+        mask_expanded = (
+            valid_mask_2d.unsqueeze(-1).expand(-1, -1, num_experts).reshape(seq_length, -1)
+        )
+
+        # Reshape and mask routing_map
+        routing_map_reshaped = routing_map.reshape(seq_length, -1)
+        routing_map_masked = routing_map_reshaped & mask_expanded
+        tokens_per_expert = routing_map_masked.sum(dim=0)
+        num_valid_tokens = valid_mask.sum()
+    else:
+        tokens_per_expert = routing_map.reshape(seq_length, -1).sum(dim=0)
+        num_valid_tokens = routing_map.shape[0]
+
+    return tokens_per_expert, num_valid_tokens
+
+
 # MOE logging
 _MOE_LAYER_WISE_LOGGING_TRACKER = {}
 
@@ -41,82 +371,83 @@ def compute_tokens_per_expert(
     seq_length: Optional[int] = None,
     bsz: Optional[int] = None,
     num_experts: Optional[int] = None,
-):
+) -> Tuple[torch.Tensor, Union[int, torch.Tensor]]:
     """Compute tokens_per_expert and num_valid_tokens with optional padding mask.
 
-    This function provides a unified way to compute token counts for MoE auxiliary losses.
-    It supports both standard aux_loss/global_aux_loss and sequence-level aux_loss.
+    This function provides a unified way to compute token counts for different aux loss types:
+    - aux_loss: Standard load balancing loss across full batch
+    - seq_aux_loss: Per-sequence load balancing loss
+    - global_aux_loss: Load balancing loss across global batch (all ranks)
 
     Padding Mask Semantics:
-        The padding_mask follows the standard masking convention where:
         - True  = padding position (should be EXCLUDED from computation)
         - False = valid position (should be INCLUDED in computation)
-        This is consistent with common attention mask conventions in transformer models.
 
-    Workflow:
-        1. If padding_mask is provided, create an inverted mask (~padding_mask) for valid tokens
-        2. Apply the valid mask to routing_map to exclude padding tokens
-        3. Sum the masked routing_map to get tokens_per_expert
-        4. Count valid tokens by summing the inverted mask
+    For aux_loss / global_aux_loss (reshape_for_seq_aux=False):
+    ----------------------------------------------------------
+    Input shapes:
+        - routing_map: [num_tokens, num_experts]
+        - padding_mask: [num_tokens]
+
+    Computation (delegated to apply_padding_mask_to_routing_map):
+        1. Invert mask: valid_mask = ~padding_mask
+        2. Expand mask to [num_tokens, 1] for broadcasting
+        3. Apply: routing_map_masked = routing_map & mask_expanded
+        4. Sum: tokens_per_expert = routing_map_masked.sum(dim=0)  # [num_experts]
+
+    For seq_aux_loss (reshape_for_seq_aux=True):
+    -------------------------------------------
+    The seq_aux_loss computes load balancing per sequence in the batch.
+    We reshape routing_map to [seq_length, bsz * num_experts] so that
+    each (batch_element, expert) pair gets its own column.
+
+    Input shapes:
+        - routing_map: [num_tokens, num_experts]
+        - padding_mask: [num_tokens]
+
+    Computation (delegated to _compute_tokens_for_seq_aux_loss):
+        1. Reshape routing_map to [seq_length, bsz * num_experts]
+        2. Reshape padding_mask to [seq_length, bsz]
+        3. Expand mask to [seq_length, bsz * num_experts]
+        4. Apply mask and sum along seq_length dimension
+        5. Result: tokens_per_expert shape [bsz * num_experts]
 
     Args:
-        routing_map (torch.Tensor): Token to expert routing map.
-            - For aux_loss/global_aux_loss: shape [num_tokens, num_experts]
-            - For seq_aux_loss: shape [num_tokens, num_experts] but will be reshaped
-        padding_mask (torch.Tensor, optional): Boolean mask indicating padding positions.
-            Shape [num_tokens]. True = padding (exclude), False = valid (include).
+        routing_map: Token to expert routing map, shape [num_tokens, num_experts].
+        padding_mask: Boolean mask [num_tokens]. True = padding (exclude), False = valid (include).
             Defaults to None (all tokens are valid).
-        reshape_for_seq_aux (bool): If True, reshape routing_map for seq_aux_loss computation.
-        seq_length (int, optional): Sequence length, required when reshape_for_seq_aux=True.
-        bsz (int, optional): Batch size, required when reshape_for_seq_aux=True.
-        num_experts (int, optional): Number of experts, required when reshape_for_seq_aux=True.
+        reshape_for_seq_aux: If True, reshape for seq_aux_loss computation.
+        seq_length: Required when reshape_for_seq_aux=True.
+        bsz: Required when reshape_for_seq_aux=True.
+        num_experts: Required when reshape_for_seq_aux=True.
 
     Returns:
-        tuple: (tokens_per_expert, num_valid_tokens)
-            - tokens_per_expert (torch.Tensor): Number of tokens per expert, shape [num_experts]
-                or [bsz * num_experts] for seq_aux_loss
-            - num_valid_tokens (torch.Tensor or int): Number of valid (non-padding) tokens
+        Tuple of (tokens_per_expert, num_valid_tokens):
+        - tokens_per_expert: [num_experts] or [bsz * num_experts] for seq_aux
+        - num_valid_tokens: Scalar count of valid tokens
+
+    Example:
+        >>> # Standard aux_loss case
+        >>> routing_map = torch.randint(0, 2, (64, 8)).bool()  # 64 tokens, 8 experts
+        >>> # Last half is padding (True = exclude)
+        >>> padding_mask = torch.zeros(64, dtype=torch.bool)
+        >>> padding_mask[32:] = True
+        >>> tokens_per_expert, num_valid = compute_tokens_per_expert(
+        ...     routing_map, padding_mask
+        ... )
+        >>> assert num_valid == 32
     """
     if reshape_for_seq_aux:
-        # seq aux loss
-        assert (
-            seq_length is not None and bsz is not None and num_experts is not None
-        ), "seq_length, bsz, and num_experts must be provided when reshape_for_seq_aux=True"
-
-        if padding_mask is not None:
-            # Invert padding_mask: True (padding) -> False, False (valid) -> True
-            valid_mask = ~padding_mask
-            # Reshape valid_mask to [seq_length, bsz]
-            valid_mask_reshaped = valid_mask.reshape(seq_length, bsz)
-            # Expand to match routing_map after reshape [seq_length, bsz * num_experts]
-            mask_expanded = (
-                valid_mask_reshaped.unsqueeze(-1)
-                .expand(-1, -1, num_experts)
-                .reshape(seq_length, -1)
+        # Validate required parameters
+        if seq_length is None or bsz is None or num_experts is None:
+            raise ValueError(
+                "seq_length, bsz, and num_experts must be provided when reshape_for_seq_aux=True"
             )
-            routing_map_masked = routing_map.reshape(seq_length, -1) & mask_expanded
-            tokens_per_expert = routing_map_masked.sum(dim=0)
-            # Count valid tokens only (where padding_mask is False)
-            num_valid_tokens = valid_mask.sum()
-        else:
-            tokens_per_expert = routing_map.reshape(seq_length, -1).sum(dim=0)
-            num_valid_tokens = routing_map.shape[0]
+        return _compute_tokens_for_seq_aux_loss(
+            routing_map, padding_mask, seq_length, bsz, num_experts
+        )
     else:
-        # aux_loss or global_aux_loss
-        if padding_mask is not None:
-            # Invert padding_mask: True (padding) -> False, False (valid) -> True
-            valid_mask = ~padding_mask
-            # routing_map: [num_tokens, num_experts], valid_mask: [num_tokens]
-            mask_expanded = valid_mask.unsqueeze(-1)
-            routing_map_masked = routing_map & mask_expanded
-            tokens_per_expert = routing_map_masked.sum(dim=0)
-            # Count valid tokens only (where padding_mask is False)
-            num_valid_tokens = valid_mask.sum()
-        else:
-            tokens_per_expert = routing_map.sum(dim=0)
-            num_valid_tokens = routing_map.shape[0]
-
-    return tokens_per_expert, num_valid_tokens
+        return apply_padding_mask_to_routing_map(routing_map, padding_mask)
 
 
 def switch_load_balancing_loss_func(
@@ -185,10 +516,7 @@ def switch_load_balancing_loss_func(
         torch.Tensor: The auxiliary loss for load balancing.
     """
     # Apply padding mask to probs if provided
-    if padding_mask is not None:
-        # Invert padding_mask: True (padding) -> 0, False (valid) -> 1
-        valid_mask = (~padding_mask).unsqueeze(-1)
-        probs = probs * valid_mask
+    probs = apply_padding_mask_to_scores(probs, padding_mask)
 
     if fused:
         if not HAVE_TE or fused_moe_aux_loss is None:
@@ -776,10 +1104,7 @@ def compute_routing_scores_for_aux_loss(
     routing_map = torch.zeros_like(logits).int().scatter(1, top_indices, 1).bool()
 
     # Apply padding mask to scores if provided
-    if padding_mask is not None:
-        # Invert padding_mask: True (padding) -> 0, False (valid) -> 1
-        valid_mask = (~padding_mask).unsqueeze(-1).to(scores.dtype)
-        scores = scores * valid_mask
+    scores = apply_padding_mask_to_scores(scores, padding_mask)
     return routing_map, scores
 
 
diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py
index 379ff4c9dd2..8aa5e315a6d 100644
--- a/megatron/core/transformer/moe/router.py
+++ b/megatron/core/transformer/moe/router.py
@@ -142,6 +142,34 @@ class TopKRouter(Router):
         scores: The scores after score function used to select the experts and calculate aux loss.
         probs: The topk weights used to combined the experts' outputs.
         routing_map: The masked routing map between tokens and experts.
+
+    Padding Mask Support:
+    ====================
+    The router accepts an optional padding_mask to exclude padding tokens from
+    auxiliary loss calculations. This is important for variable-length sequences
+    where shorter sequences are padded.
+
+    When padding_mask is provided:
+    1. Z-loss is computed only over valid tokens (mean excludes padding)
+    2. Load balancing loss counts only valid token assignments
+    3. Expert bias updates exclude padding token routings
+
+    The mask does NOT affect the actual routing decisions - padding tokens are
+    still routed to experts, but they don't contribute to loss calculations.
+    This design choice ensures:
+    - Consistent output shapes regardless of padding
+    - No changes to inference behavior
+    - Only training loss computation is affected
+
+    Padding Mask Semantics (IMPORTANT):
+    -----------------------------------
+    - True  = padding position (should be EXCLUDED from computation)
+    - False = valid position (should be INCLUDED in computation)
+
+    Padding Mask Shape Convention:
+    -----------------------------
+    - Input to forward(): [seq_length, batch_size] (matches hidden_states layout)
+    - Internal usage: [num_tokens] (flattened for per-token operations)
     """
 
     def __init__(
@@ -598,8 +626,8 @@ def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = N
         Args:
             logits (torch.Tensor): Logits tensor after gating.
             padding_mask (torch.Tensor, optional): Boolean mask indicating padding positions.
-                                                   Shape = [seq_length, bsz]. True = padding,
-                                                   False = valid (include). Defaults to None.
+                                                   Shape = [seq_length, bsz]. True=padding(exclude),
+                                                   False=valid(include). Defaults to None.
 
         Returns:
             probs (torch.Tensor): The probabilities of token to experts assignment.
@@ -688,8 +716,8 @@ def forward(self, input: torch.Tensor, padding_mask: Optional[torch.Tensor] = No
         Args:
             input (torch.Tensor): Input tensor.
             padding_mask (torch.Tensor, optional): Boolean mask indicating padding positions.
-                                                   Shape = [seq_length, bsz]. True = padding,
-                                                   False = valid (include). Defaults to None.
+                                                   Shape = [seq_length, bsz]. True=padding(exclude),
+                                                   False=valid(include). Defaults to None.
         """
         self._maintain_float32_expert_bias()
 
diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py
index 5ebb3753422..93be17e9326 100644
--- a/megatron/core/transformer/transformer_layer.py
+++ b/megatron/core/transformer/transformer_layer.py
@@ -597,12 +597,27 @@ def _forward_mlp(self, hidden_states, inference_context=None, padding_mask=None)
         """
         Perform a forward pass through the feed-forward layer.
 
+        For MoE layers, this method handles padding mask forwarding to ensure
+        padding tokens are excluded from auxiliary loss calculations.
+
+        Padding Mask Flow:
+            1. User provides padding_mask to GPTModel.forward() with shape [bsz, seq_length]
+            2. TransformerBlock passes it through to each TransformerLayer
+            3. This method passes it to MoELayer.forward() (for MoE layers only)
+            4. MoELayer uses PaddingMaskInfo to transform shape and handle masking
+
+        Padding Mask Semantics:
+            - True  = padding position (should be EXCLUDED from computation)
+            - False = valid position (should be INCLUDED in computation)
+
         Args:
             hidden_states (Tensor): Transformed hidden states before the MLP layernorm.
+                Shape [seq_length, batch_size, hidden_size].
             inference_context: Inference context for optimizations.
             padding_mask (Tensor, optional): Padding mask for MoE routing.
                 Shape [bsz, seq_length]. True = padding (exclude), False = valid (include).
-                Only used for MoE layers to exclude padding tokens from routing computations.
+                Only used for MoE layers to exclude padding tokens from aux loss computations.
+                The MoELayer will internally transform this to [seq_length, bsz] format.
 
         Returns:
             output (Tensor): Transformed hidden states of shape [s, b, h].

From 621874136c31b5cfb1763837ae3faef315caa742 Mon Sep 17 00:00:00 2001
From: HaochenYuan <haocheny@nvidia.com>
Date: Thu, 11 Dec 2025 01:11:01 -0800
Subject: [PATCH 11/32] refactor

---
 megatron/core/transformer/mlp.py              |   2 +-
 megatron/core/transformer/moe/moe_layer.py    |  38 +-
 megatron/core/transformer/moe/moe_utils.py    | 471 ++----------------
 megatron/core/transformer/moe/router.py       | 156 ++----
 .../core/transformer/transformer_layer.py     |  42 +-
 5 files changed, 79 insertions(+), 630 deletions(-)

diff --git a/megatron/core/transformer/mlp.py b/megatron/core/transformer/mlp.py
index 9602beb2f71..eb78cc0fde2 100644
--- a/megatron/core/transformer/mlp.py
+++ b/megatron/core/transformer/mlp.py
@@ -137,7 +137,7 @@ def __init__(
             tp_group=tp_group,
         )
 
-    def forward(self, hidden_states, per_token_scale=None):
+    def forward(self, hidden_states, per_token_scale=None, **kwargs):
         """Perform the forward pass through the MLP block."""
         # [s, b, 4 * h/p]
         nvtx_range_push(suffix="linear_fc1")
diff --git a/megatron/core/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py
index bca90c7f916..ca7ead48540 100644
--- a/megatron/core/transformer/moe/moe_layer.py
+++ b/megatron/core/transformer/moe/moe_layer.py
@@ -290,32 +290,10 @@ def forward(self, hidden_states: torch.Tensor, padding_mask: Optional[torch.Tens
         3. Expert Computation: Experts process the dispatched tokens.
         4. Combine: The outputs from the experts are combined and returned.
 
-        Padding Mask Handling:
-        =====================
-        The padding_mask indicates which positions should be excluded from MoE routing
-        and auxiliary loss computation. This is useful when training with variable-length
-        sequences that are padded to a fixed length.
-
-        Padding Mask Semantics:
-            - True  = padding position (should be EXCLUDED from computation)
-            - False = valid position (should be INCLUDED in computation)
-
-        Shape Transformation:
-            - Input padding_mask from GPTModel/TransformerLayer: [bsz, seq_length]
-            - This method uses PaddingMaskInfo to transform to: [seq_length, bsz]
-            - This aligns with hidden_states shape [seq_length, bsz, hidden_size]
-
-        The mask is passed through to the router, which uses utility functions to:
-            - Exclude padding tokens from auxiliary loss calculations
-            - Exclude padding tokens from expert bias updates
-            - Exclude padding tokens from z-loss computation
-
-        Note: The mask does NOT affect actual routing decisions - padding tokens are
-        still routed to experts, but they don't contribute to loss calculations.
-
         Args:
             hidden_states (torch.Tensor): The input tensor shape [seq_length, bsz, hidden_size].
             padding_mask (torch.Tensor, optional): Boolean mask indicating padding positions.
+                used for correct auxiliary loss computation for packed sequence.
                 Shape = [bsz, seq_length]. True = padding (exclude), False = valid (include).
                 Defaults to None (all tokens are valid).
 
@@ -328,11 +306,9 @@ def forward(self, hidden_states: torch.Tensor, padding_mask: Optional[torch.Tens
                 "are enabled without also enabling sequence parallelism."
             )
 
-        # Use PaddingMaskInfo for unified mask handling
-        # Transform from [bsz, seq_length] to [seq_length, bsz] format
-        seq_length, bsz = hidden_states.shape[:2]
-        mask_info = PaddingMaskInfo.from_bsz_seq(padding_mask, seq_length, bsz)
-        padding_mask_seq_bsz = mask_info.seq_bsz
+        # Transpose from [bsz, seq_length] to [seq_length, bsz] to align with hidden_states
+        if padding_mask is not None:
+            padding_mask = padding_mask.transpose(0, 1).bool()
 
         # MoE forward: route -> dispatch -> compute -> combine
         def custom_forward(hidden_states, padding_mask=None):
@@ -361,14 +337,14 @@ def custom_forward(hidden_states, padding_mask=None):
                     tensor_parallel.random.get_cuda_rng_tracker,
                     parallel_state.get_tensor_model_parallel_group(),
                     hidden_states,
-                    padding_mask_seq_bsz,
+                    padding_mask,
                 )
             else:
                 outputs = tensor_parallel.checkpoint(
-                    custom_forward, False, hidden_states, padding_mask_seq_bsz
+                    custom_forward, False, hidden_states, padding_mask
                 )
         else:
-            outputs = custom_forward(hidden_states, padding_mask_seq_bsz)
+            outputs = custom_forward(hidden_states, padding_mask)
 
         return outputs
 
diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py
index aa3b15d9c79..6c1342bda80 100644
--- a/megatron/core/transformer/moe/moe_utils.py
+++ b/megatron/core/transformer/moe/moe_utils.py
@@ -1,50 +1,8 @@
-# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-
-"""
-MoE Utility Functions
-=====================
-
-This module provides utilities for Mixture of Experts (MoE) routing and loss computation.
-
-Padding Mask Handling
----------------------
-MoE models need special handling for padding tokens in variable-length sequences.
-Without proper masking, padding tokens would:
-1. Artificially inflate expert load statistics
-2. Contribute to z-loss despite being meaningless
-3. Skew load balancing metrics
-
-This module provides unified utilities for padding mask handling:
-
-Padding Mask Semantics (IMPORTANT):
-    The padding_mask follows the standard masking convention where:
-    - True  = padding position (should be EXCLUDED from computation)
-    - False = valid position (should be INCLUDED in computation)
-    This is consistent with common attention mask conventions in transformer models.
-
-Shape Conventions:
-    - User input (GPTModel): [batch_size, seq_length]
-    - Internal format: [seq_length, batch_size] (Megatron convention)
-    - Flattened: [num_tokens] where num_tokens = seq_length * batch_size
-
-Key Classes and Functions:
-    - PaddingMaskInfo: Unified container for padding mask with different shape views
-    - apply_padding_mask_to_routing_map(): Apply mask to routing decisions
-    - apply_padding_mask_to_scores(): Zero out scores for padding tokens
-    - compute_tokens_per_expert(): Unified token counting with mask support
-
-Example Workflow:
-    1. User provides padding_mask [bsz, seq] to GPTModel.forward()
-    2. MoELayer transposes to [seq, bsz] to match hidden_states
-    3. Router creates PaddingMaskInfo for unified handling
-    4. Router flattens mask to [num_tokens] for per-token operations
-    5. Aux loss functions use mask utilities to exclude padding from statistics
-"""
-
 import math
 from dataclasses import dataclass
 from typing import List, Optional, Union
 
+from megatron.core.tensor_parallel.mappings import reduce_from_tensor_model_parallel_region
 import torch
 
 from megatron.core import parallel_state
@@ -71,385 +29,10 @@
     HAVE_TE = False
 
 
-from dataclasses import dataclass
-from typing import Tuple
-
-
-@dataclass
-class PaddingMaskInfo:
-    """Container for padding mask with different views/shapes.
-
-    This class provides a unified interface for handling padding masks throughout
-    the MoE routing and auxiliary loss computation pipeline.
-
-    Padding Mask Semantics:
-        - True  = padding position (should be EXCLUDED from computation)
-        - False = valid position (should be INCLUDED in computation)
-
-    The padding mask workflow in MoE:
-    ================================
-    1. User provides padding_mask to GPTModel.forward() with shape [batch_size, seq_length]
-       - True = padding token (exclude), False = valid token (include)
-
-    2. MoELayer transposes it to [seq_length, batch_size] to match hidden_states layout
-       (Megatron uses [seq, batch, hidden] format internally)
-
-    3. Router flattens it to [num_tokens] where num_tokens = seq_length * batch_size
-
-    4. For operations with routing_map [num_tokens, num_experts], we use utility
-       functions that handle mask expansion internally
-
-    Shape Conventions:
-    -----------------
-    - bsz_seq: [batch_size, seq_length] - User input format
-    - seq_bsz: [seq_length, batch_size] - Internal Megatron format
-    - flat: [num_tokens] - Flattened for per-token operations
-
-    Attributes:
-        seq_bsz: Mask in [seq_length, batch_size] format (None if no mask provided)
-        flat: Flattened mask in [num_tokens] format
-        num_valid_tokens: Count of valid (non-padding) tokens
-        seq_length: Sequence length
-        batch_size: Batch size
-
-    Example Usage:
-    -------------
-    >>> # Create from user-provided mask [bsz, seq]
-    >>> mask_info = PaddingMaskInfo.from_bsz_seq(padding_mask, seq_length=32, batch_size=4)
-    >>> # Check if mask exists
-    >>> if mask_info.has_mask:
-    ...     # Use flattened mask in aux loss computation
-    ...     tokens_per_expert, num_valid = apply_padding_mask_to_routing_map(
-    ...         routing_map, mask_info.flat
-    ...     )
-    """
-
-    # Mask in [seq_length, batch_size] format (None if no mask provided)
-    seq_bsz: Optional[torch.Tensor]
-
-    # Flattened mask in [num_tokens] format
-    flat: Optional[torch.Tensor]
-
-    # Number of valid (non-padding) tokens
-    num_valid_tokens: Optional[torch.Tensor]
-
-    # Metadata
-    seq_length: int
-    batch_size: int
-
-    @classmethod
-    def from_bsz_seq(
-        cls, padding_mask: Optional[torch.Tensor], seq_length: int, batch_size: int
-    ) -> "PaddingMaskInfo":
-        """Create PaddingMaskInfo from user-provided mask in [batch_size, seq_length] format.
-
-        Args:
-            padding_mask: Boolean tensor [batch_size, seq_length].
-                         True = padding (exclude), False = valid (include).
-                         If None, creates an "all valid" mask info.
-            seq_length: Sequence length.
-            batch_size: Batch size.
-
-        Returns:
-            PaddingMaskInfo with all shape variants populated.
-        """
-        if padding_mask is None:
-            return cls(
-                seq_bsz=None,
-                flat=None,
-                num_valid_tokens=None,
-                seq_length=seq_length,
-                batch_size=batch_size,
-            )
-
-        # Transpose to [seq_length, batch_size] for internal use
-        seq_bsz = padding_mask.transpose(0, 1).bool()
-
-        # Flatten to [num_tokens]
-        flat = seq_bsz.reshape(-1)
-
-        # Count valid tokens (where padding_mask is False)
-        num_valid_tokens = (~flat).sum()
-
-        return cls(
-            seq_bsz=seq_bsz,
-            flat=flat,
-            num_valid_tokens=num_valid_tokens,
-            seq_length=seq_length,
-            batch_size=batch_size,
-        )
-
-    @classmethod
-    def from_seq_bsz(
-        cls, padding_mask: Optional[torch.Tensor], seq_length: int, batch_size: int
-    ) -> "PaddingMaskInfo":
-        """Create PaddingMaskInfo from mask already in [seq_length, batch_size] format.
-
-        Use this when the mask has already been transposed (e.g., inside MoELayer).
-
-        Args:
-            padding_mask: Boolean tensor [seq_length, batch_size].
-                         True = padding (exclude), False = valid (include).
-                         If None, creates an "all valid" mask info.
-            seq_length: Sequence length.
-            batch_size: Batch size.
-
-        Returns:
-            PaddingMaskInfo with all shape variants populated.
-        """
-        if padding_mask is None:
-            return cls(
-                seq_bsz=None,
-                flat=None,
-                num_valid_tokens=None,
-                seq_length=seq_length,
-                batch_size=batch_size,
-            )
-
-        seq_bsz = padding_mask.bool()
-        flat = seq_bsz.reshape(-1)
-        # Count valid tokens (where padding_mask is False)
-        num_valid_tokens = (~flat).sum()
-
-        return cls(
-            seq_bsz=seq_bsz,
-            flat=flat,
-            num_valid_tokens=num_valid_tokens,
-            seq_length=seq_length,
-            batch_size=batch_size,
-        )
-
-    @property
-    def has_mask(self) -> bool:
-        """Check if a padding mask was provided."""
-        return self.flat is not None
-
-    @property
-    def num_total_tokens(self) -> int:
-        """Total number of tokens (including padding)."""
-        return self.seq_length * self.batch_size
-
-
-def apply_padding_mask_to_routing_map(
-    routing_map: torch.Tensor, padding_mask_flat: Optional[torch.Tensor]
-) -> Tuple[torch.Tensor, Union[int, torch.Tensor]]:
-    """Apply padding mask to routing map and compute tokens per expert.
-
-    This is a unified function to handle the common pattern of:
-    1. Inverting padding_mask (True=padding -> False, False=valid -> True)
-    2. Expanding the valid mask from [num_tokens] to [num_tokens, 1]
-    3. Masking routing_map with the expanded mask
-    4. Computing tokens_per_expert from the masked routing_map
-
-    Args:
-        routing_map: Boolean tensor [num_tokens, num_experts] indicating token-expert assignments.
-        padding_mask_flat: Boolean tensor [num_tokens].
-            True = padding (exclude), False = valid (include).
-            If None, no masking is applied.
-
-    Returns:
-        Tuple of:
-        - tokens_per_expert: Tensor [num_experts] with count of valid tokens per expert
-        - num_valid_tokens: Scalar tensor or int with total valid token count
-
-    Example:
-        >>> routing_map = torch.tensor([[True, False], [False, True], [True, False]])
-        >>> # Last token is padding (True = exclude)
-        >>> padding_mask = torch.tensor([False, False, True])
-        >>> tokens_per_expert, num_valid
-            = apply_padding_mask_to_routing_map(routing_map, padding_mask)
-        >>> print(tokens_per_expert)  # tensor([1, 1]) - only counts first two tokens
-        >>> print(num_valid)  # tensor(2)
-    """
-    if padding_mask_flat is None:
-        tokens_per_expert = routing_map.sum(dim=0)
-        num_valid_tokens = routing_map.shape[0]
-    else:
-        # Invert: True (padding) -> False, False (valid) -> True
-        valid_mask = ~padding_mask_flat
-        # Expand mask for broadcasting: [num_tokens] -> [num_tokens, 1]
-        mask_expanded = valid_mask.unsqueeze(-1)
-        # Apply mask: only count tokens where valid_mask is True
-        routing_map_masked = routing_map & mask_expanded
-        tokens_per_expert = routing_map_masked.sum(dim=0)
-        num_valid_tokens = valid_mask.sum()
-
-    return tokens_per_expert, num_valid_tokens
-
-
-def apply_padding_mask_to_scores(
-    scores: torch.Tensor, padding_mask_flat: Optional[torch.Tensor]
-) -> torch.Tensor:
-    """Zero out scores for padding tokens.
-
-    Args:
-        scores: Tensor [num_tokens, num_experts] with routing scores/probabilities.
-        padding_mask_flat: Boolean tensor [num_tokens].
-            True = padding (exclude), False = valid (include).
-            If None, no masking is applied.
-
-    Returns:
-        Masked scores with padding positions set to 0.
-
-    Example:
-        >>> scores = torch.tensor([[0.5, 0.5], [0.3, 0.7], [0.8, 0.2]])
-        >>> # Last token is padding (True = exclude)
-        >>> padding_mask = torch.tensor([False, False, True])
-        >>> masked_scores = apply_padding_mask_to_scores(scores, padding_mask)
-        >>> print(masked_scores)  # Last row is zeroed out
-    """
-    if padding_mask_flat is None:
-        return scores
-
-    # Invert: True (padding) -> 0, False (valid) -> 1
-    valid_mask = (~padding_mask_flat).unsqueeze(-1).to(scores.dtype)
-    return scores * valid_mask
-
-
-def _compute_tokens_for_seq_aux_loss(
-    routing_map: torch.Tensor,
-    padding_mask: Optional[torch.Tensor],
-    seq_length: int,
-    bsz: int,
-    num_experts: int,
-) -> Tuple[torch.Tensor, Union[int, torch.Tensor]]:
-    """Internal helper for seq_aux_loss token counting.
-
-    For seq_aux_loss, we need per-sequence, per-expert token counts.
-    This requires reshaping the data to separate batch elements.
-
-    Reshape logic:
-    - routing_map [num_tokens, num_experts] -> [seq_length, bsz, num_experts]
-                                            -> [seq_length, bsz * num_experts]
-    - padding_mask [num_tokens] -> [seq_length, bsz] -> broadcast to [seq_length, bsz * num_experts]
-
-    Args:
-        routing_map: Token to expert routing map [num_tokens, num_experts].
-        padding_mask: Boolean mask [num_tokens]. True = padding (exclude), False = valid (include).
-        seq_length: Sequence length.
-        bsz: Batch size.
-        num_experts: Number of experts.
-
-    Returns:
-        Tuple of (tokens_per_expert, num_valid_tokens):
-        - tokens_per_expert: [bsz * num_experts]
-        - num_valid_tokens: Scalar count of valid tokens
-    """
-    if padding_mask is not None:
-        # Invert: True (padding) -> False, False (valid) -> True
-        valid_mask = ~padding_mask
-        # Reshape valid_mask: [num_tokens] -> [seq_length, bsz]
-        valid_mask_2d = valid_mask.reshape(seq_length, bsz)
-
-        # Expand for broadcasting with routing_map
-        # [seq_length, bsz] -> [seq_length, bsz, 1] -> [seq_length, bsz, num_experts]
-        # -> [seq_length, bsz * num_experts]
-        mask_expanded = (
-            valid_mask_2d.unsqueeze(-1).expand(-1, -1, num_experts).reshape(seq_length, -1)
-        )
-
-        # Reshape and mask routing_map
-        routing_map_reshaped = routing_map.reshape(seq_length, -1)
-        routing_map_masked = routing_map_reshaped & mask_expanded
-        tokens_per_expert = routing_map_masked.sum(dim=0)
-        num_valid_tokens = valid_mask.sum()
-    else:
-        tokens_per_expert = routing_map.reshape(seq_length, -1).sum(dim=0)
-        num_valid_tokens = routing_map.shape[0]
-
-    return tokens_per_expert, num_valid_tokens
-
-
 # MOE logging
 _MOE_LAYER_WISE_LOGGING_TRACKER = {}
 
 
-def compute_tokens_per_expert(
-    routing_map: torch.Tensor,
-    padding_mask: Optional[torch.Tensor] = None,
-    reshape_for_seq_aux: bool = False,
-    seq_length: Optional[int] = None,
-    bsz: Optional[int] = None,
-    num_experts: Optional[int] = None,
-) -> Tuple[torch.Tensor, Union[int, torch.Tensor]]:
-    """Compute tokens_per_expert and num_valid_tokens with optional padding mask.
-
-    This function provides a unified way to compute token counts for different aux loss types:
-    - aux_loss: Standard load balancing loss across full batch
-    - seq_aux_loss: Per-sequence load balancing loss
-    - global_aux_loss: Load balancing loss across global batch (all ranks)
-
-    Padding Mask Semantics:
-        - True  = padding position (should be EXCLUDED from computation)
-        - False = valid position (should be INCLUDED in computation)
-
-    For aux_loss / global_aux_loss (reshape_for_seq_aux=False):
-    ----------------------------------------------------------
-    Input shapes:
-        - routing_map: [num_tokens, num_experts]
-        - padding_mask: [num_tokens]
-
-    Computation (delegated to apply_padding_mask_to_routing_map):
-        1. Invert mask: valid_mask = ~padding_mask
-        2. Expand mask to [num_tokens, 1] for broadcasting
-        3. Apply: routing_map_masked = routing_map & mask_expanded
-        4. Sum: tokens_per_expert = routing_map_masked.sum(dim=0)  # [num_experts]
-
-    For seq_aux_loss (reshape_for_seq_aux=True):
-    -------------------------------------------
-    The seq_aux_loss computes load balancing per sequence in the batch.
-    We reshape routing_map to [seq_length, bsz * num_experts] so that
-    each (batch_element, expert) pair gets its own column.
-
-    Input shapes:
-        - routing_map: [num_tokens, num_experts]
-        - padding_mask: [num_tokens]
-
-    Computation (delegated to _compute_tokens_for_seq_aux_loss):
-        1. Reshape routing_map to [seq_length, bsz * num_experts]
-        2. Reshape padding_mask to [seq_length, bsz]
-        3. Expand mask to [seq_length, bsz * num_experts]
-        4. Apply mask and sum along seq_length dimension
-        5. Result: tokens_per_expert shape [bsz * num_experts]
-
-    Args:
-        routing_map: Token to expert routing map, shape [num_tokens, num_experts].
-        padding_mask: Boolean mask [num_tokens]. True = padding (exclude), False = valid (include).
-            Defaults to None (all tokens are valid).
-        reshape_for_seq_aux: If True, reshape for seq_aux_loss computation.
-        seq_length: Required when reshape_for_seq_aux=True.
-        bsz: Required when reshape_for_seq_aux=True.
-        num_experts: Required when reshape_for_seq_aux=True.
-
-    Returns:
-        Tuple of (tokens_per_expert, num_valid_tokens):
-        - tokens_per_expert: [num_experts] or [bsz * num_experts] for seq_aux
-        - num_valid_tokens: Scalar count of valid tokens
-
-    Example:
-        >>> # Standard aux_loss case
-        >>> routing_map = torch.randint(0, 2, (64, 8)).bool()  # 64 tokens, 8 experts
-        >>> # Last half is padding (True = exclude)
-        >>> padding_mask = torch.zeros(64, dtype=torch.bool)
-        >>> padding_mask[32:] = True
-        >>> tokens_per_expert, num_valid = compute_tokens_per_expert(
-        ...     routing_map, padding_mask
-        ... )
-        >>> assert num_valid == 32
-    """
-    if reshape_for_seq_aux:
-        # Validate required parameters
-        if seq_length is None or bsz is None or num_experts is None:
-            raise ValueError(
-                "seq_length, bsz, and num_experts must be provided when reshape_for_seq_aux=True"
-            )
-        return _compute_tokens_for_seq_aux_loss(
-            routing_map, padding_mask, seq_length, bsz, num_experts
-        )
-    else:
-        return apply_padding_mask_to_routing_map(routing_map, padding_mask)
-
-
 def switch_load_balancing_loss_func(
     probs: torch.Tensor,
     tokens_per_expert: torch.Tensor,
@@ -458,7 +41,6 @@ def switch_load_balancing_loss_func(
     num_experts: int,
     moe_aux_loss_coeff: float,
     fused: bool = False,
-    padding_mask: Optional[torch.Tensor] = None,
 ):
     """Calculate the auxiliary loss for load balancing.
     Refer to the Switch Transformer (https://arxiv.org/abs/2101.03961)
@@ -515,9 +97,6 @@ def switch_load_balancing_loss_func(
     Returns:
         torch.Tensor: The auxiliary loss for load balancing.
     """
-    # Apply padding mask to probs if provided
-    probs = apply_padding_mask_to_scores(probs, padding_mask)
-
     if fused:
         if not HAVE_TE or fused_moe_aux_loss is None:
             raise ValueError("fused_moe_aux_loss is not available. Please install TE >= 2.7.0.")
@@ -603,6 +182,23 @@ def get_capacity(num_tokens: int, num_experts: int, capacity_factor: float, min_
         capacity = min_capacity
     return capacity
 
+def get_tokens_per_expert_and_token_count(
+    routing_map: torch.Tensor, 
+    reduce_group: torch.distributed.ProcessGroup, 
+    topk: int = None,
+    with_padding_mask: bool = False
+) -> torch.Tensor:
+    local_tokens_per_expert = routing_map.sum(dim=0)
+    global_tokens_per_expert = reduce_from_tensor_model_parallel_region(
+        local_tokens_per_expert, reduce_group
+    )
+    if with_padding_mask:
+        local_num_tokens = local_tokens_per_expert.sum() / topk
+        total_num_tokens = global_tokens_per_expert.sum() / topk
+    else:
+        local_num_tokens = routing_map.shape[0]
+        total_num_tokens = local_num_tokens * reduce_group.size()
+    return global_tokens_per_expert, local_num_tokens, total_num_tokens
 
 class MoEAuxLossAutoScaler(torch.autograd.Function):
     """An AutoScaler that triggers the backward pass and scales the grad for auxiliary loss."""
@@ -1083,28 +679,27 @@ def compute_routing_scores_for_aux_loss(
             raise ValueError(
                 "fused_compute_score_for_moe_aux_loss is not available. Please install TE >= 2.6.0."
             )
-        # Note: fused implementation does not support padding_mask yet
-        if padding_mask is not None:
-            raise ValueError(
-                "Fused compute_routing_scores does not support padding_mask. Set fused=False."
-            )
-        return fused_compute_score_for_moe_aux_loss(
+        routing_map, scores = fused_compute_score_for_moe_aux_loss(
             logits=logits, topk=topk, score_function=score_function
         )
-
-    if score_function == "softmax":
-        scores = torch.softmax(logits, dim=-1, dtype=torch.float32)
-    elif score_function == "sigmoid":
-        scores = torch.sigmoid(logits)
-        scores = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20)
     else:
-        raise ValueError(f"Invalid score_function: {score_function}")
+        if score_function == "softmax":
+            scores = torch.softmax(logits, dim=-1, dtype=torch.float32)
+        elif score_function == "sigmoid":
+            scores = torch.sigmoid(logits)
+            scores = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20)
+        else:
+            raise ValueError(f"Invalid score_function: {score_function}")
 
-    _, top_indices = torch.topk(scores, k=topk, dim=1)
-    routing_map = torch.zeros_like(logits).int().scatter(1, top_indices, 1).bool()
+        _, top_indices = torch.topk(scores, k=topk, dim=1)
+        routing_map = torch.zeros_like(logits).int().scatter(1, top_indices, 1).bool()
 
     # Apply padding mask to scores if provided
-    scores = apply_padding_mask_to_scores(scores, padding_mask)
+    if padding_mask is not None:
+        # Invert padding_mask and make True indicates valid tokens
+        valid_mask = (~padding_mask).unsqueeze(-1)
+        routing_map = routing_map * valid_mask
+        scores = scores * valid_mask
     return routing_map, scores
 
 
diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py
index 8aa5e315a6d..7dfd780a164 100644
--- a/megatron/core/transformer/moe/router.py
+++ b/megatron/core/transformer/moe/router.py
@@ -14,7 +14,7 @@
     apply_random_logits,
     apply_router_token_dropping,
     compute_routing_scores_for_aux_loss,
-    compute_tokens_per_expert,
+    get_tokens_per_expert_and_token_count,
     router_gating_linear,
     save_to_aux_losses_tracker,
     sinkhorn,
@@ -142,34 +142,6 @@ class TopKRouter(Router):
         scores: The scores after score function used to select the experts and calculate aux loss.
         probs: The topk weights used to combined the experts' outputs.
         routing_map: The masked routing map between tokens and experts.
-
-    Padding Mask Support:
-    ====================
-    The router accepts an optional padding_mask to exclude padding tokens from
-    auxiliary loss calculations. This is important for variable-length sequences
-    where shorter sequences are padded.
-
-    When padding_mask is provided:
-    1. Z-loss is computed only over valid tokens (mean excludes padding)
-    2. Load balancing loss counts only valid token assignments
-    3. Expert bias updates exclude padding token routings
-
-    The mask does NOT affect the actual routing decisions - padding tokens are
-    still routed to experts, but they don't contribute to loss calculations.
-    This design choice ensures:
-    - Consistent output shapes regardless of padding
-    - No changes to inference behavior
-    - Only training loss computation is affected
-
-    Padding Mask Semantics (IMPORTANT):
-    -----------------------------------
-    - True  = padding position (should be EXCLUDED from computation)
-    - False = valid position (should be INCLUDED in computation)
-
-    Padding Mask Shape Convention:
-    -----------------------------
-    - Input to forward(): [seq_length, batch_size] (matches hidden_states layout)
-    - Internal usage: [num_tokens] (flattened for per-token operations)
     """
 
     def __init__(
@@ -301,44 +273,27 @@ def _apply_aux_loss(
         probs: torch.Tensor,
         scores_for_aux_loss: torch.Tensor,
         routing_map: torch.Tensor,
-        padding_mask: Optional[torch.Tensor] = None,
+        with_padding_mask: bool = False,
     ):
         """Apply the auxiliary loss for the given scores and routing map."""
         aux_loss_coeff = self.get_aux_loss_coeff("aux_loss")
         if aux_loss_coeff == 0:
             return probs
 
-        # Use unified function to compute tokens_per_expert and num_tokens
-        tokens_per_expert, num_tokens = compute_tokens_per_expert(
-            routing_map=routing_map, padding_mask=padding_mask, reshape_for_seq_aux=False
-        )
-
-        tokens_per_expert = reduce_from_tensor_model_parallel_region(
-            tokens_per_expert, self.tp_cp_group
+        global_tokens_per_expert, local_num_tokens, total_num_tokens = get_tokens_per_expert_and_token_count(
+            routing_map=routing_map,
+            reduce_group=self.tp_cp_group,
+            topk=self.topk,
+            with_padding_mask=with_padding_mask,
         )
-        # When padding_mask is not None, we need the global count of valid tokens across all ranks.
-        # Naively, this would require an extra all_reduce on num_valid_tokens.
-        #
-        # Optimization: We derive total_num_tokens from the already all-reduced tokens_per_expert:
-        #   - tokens_per_expert.sum() = total valid token-expert assignments (global)
-        #   - Each valid token is assigned to exactly `topk` experts
-        #   - Therefore: total_valid_tokens = tokens_per_expert.sum() / topk
-        #
-        # This avoids an extra all_reduce operation.
-        if padding_mask is not None:
-            total_num_tokens = tokens_per_expert.sum() / self.topk
-        else:
-            total_num_tokens = num_tokens * self.tp_cp_group.size()
-
         aux_loss = switch_load_balancing_loss_func(
             probs=scores_for_aux_loss,
-            tokens_per_expert=tokens_per_expert,
+            tokens_per_expert=global_tokens_per_expert,
             total_num_tokens=total_num_tokens,
             topk=self.topk,
             num_experts=self.config.num_moe_experts,
             moe_aux_loss_coeff=aux_loss_coeff,
             fused=self.config.moe_router_fusion,
-            padding_mask=padding_mask,
         )
         probs = self.attach_and_log_load_balancing_loss(
             probs,
@@ -346,7 +301,7 @@ def _apply_aux_loss(
             aux_loss,
             "load_balancing_loss",
             self.tp_cp_group,
-            valid_token_count=num_tokens,
+            valid_token_count=local_num_tokens,
         )
         return probs
 
@@ -357,7 +312,7 @@ def _apply_seq_aux_loss(
         routing_map: torch.Tensor,
         seq_length: int,
         bsz: int,
-        padding_mask: Optional[torch.Tensor] = None,
+        with_padding_mask: bool = False,
     ):
         """Apply the sequence-level auxiliary loss for the given scores and routing map.
 
@@ -372,49 +327,26 @@ def _apply_seq_aux_loss(
 
         scores_for_aux_loss = scores_for_aux_loss.reshape(seq_length, -1)
 
-        # Use unified function to compute tokens_per_expert and num_tokens
-        tokens_per_expert, num_tokens = compute_tokens_per_expert(
+        global_tokens_per_expert, local_num_tokens, total_num_tokens = get_tokens_per_expert_and_token_count(
             routing_map=routing_map,
-            padding_mask=padding_mask,
-            reshape_for_seq_aux=True,
-            seq_length=seq_length,
-            bsz=bsz,
-            num_experts=self.config.num_moe_experts,
-        )
-
-        tokens_per_expert = reduce_from_tensor_model_parallel_region(
-            tokens_per_expert, self.tp_cp_group
+            reduce_group=self.tp_cp_group,
+            with_padding_mask=with_padding_mask,
+            topk=self.topk,
         )
-        # For seq_aux_loss, total_num_tokens should be seq_length (per sequence),
-        # not seq_length * batch_size.
-        # When padding_mask is not None, compute from tokens_per_expert to get accurate count.
-        # tokens_per_expert shape: [bsz * num_experts]
-        # tokens_per_expert.sum() = seq_length * bsz * topk
-        # Divide by (bsz * topk) to get seq_length
-        if padding_mask is not None:
-            total_num_tokens = tokens_per_expert.sum() / (bsz * self.topk)
-        else:
-            total_num_tokens = seq_length * self.tp_cp_group.size()
 
         aux_loss = (
             switch_load_balancing_loss_func(
                 probs=scores_for_aux_loss,
-                tokens_per_expert=tokens_per_expert,
+                tokens_per_expert=global_tokens_per_expert,
                 total_num_tokens=total_num_tokens,
                 topk=self.topk,
                 num_experts=self.config.num_moe_experts,
                 moe_aux_loss_coeff=seq_aux_loss_coeff,
                 fused=self.config.moe_router_fusion,
-                padding_mask=None,
             )
             / bsz
         )
-        # Calculate valid token count: number of non-padding tokens
-        if padding_mask is not None:
-            # Invert padding_mask: count positions where padding_mask is False
-            valid_token_count = (~padding_mask).sum()
-        else:
-            valid_token_count = seq_length * bsz
+
 
         probs = self.attach_and_log_load_balancing_loss(
             probs,
@@ -422,7 +354,7 @@ def _apply_seq_aux_loss(
             aux_loss,
             "seq_load_balancing_loss",
             self.tp_cp_group,
-            valid_token_count=valid_token_count,
+            valid_token_count=local_num_tokens,
         )
         return probs
 
@@ -431,7 +363,7 @@ def _apply_global_aux_loss(
         probs: torch.Tensor,
         scores_for_aux_loss: torch.Tensor,
         routing_map: torch.Tensor,
-        padding_mask: Optional[torch.Tensor] = None,
+        with_padding_mask: bool = False,
     ):
         """Apply the global auxiliary loss for the given scores and routing map."""
         global_aux_loss_coeff = self.get_aux_loss_coeff("global_aux_loss")
@@ -439,31 +371,16 @@ def _apply_global_aux_loss(
             return probs
 
         # Use unified function to compute tokens_per_expert and num_tokens
-        tokens_per_expert, num_tokens = compute_tokens_per_expert(
-            routing_map=routing_map, padding_mask=padding_mask, reshape_for_seq_aux=False
-        )
-
-        tokens_per_expert = reduce_from_tensor_model_parallel_region(
-            tokens_per_expert, self.tp_dp_cp_group
+        global_tokens_per_expert, local_num_tokens, total_num_tokens = get_tokens_per_expert_and_token_count(
+            routing_map=routing_map,
+            reduce_group=self.tp_dp_cp_group,
+            with_padding_mask=with_padding_mask,
+            topk=self.topk,
         )
-
-        self.global_tokens_per_expert += tokens_per_expert
+        self.global_tokens_per_expert += global_tokens_per_expert
         self.ga_steps += 1
         averated_tokens_per_expert = self.global_tokens_per_expert / self.ga_steps
 
-        # When padding_mask is not None, we need the global count of valid tokens across all ranks.
-        # Naively, this would require an extra all_reduce on num_valid_tokens.
-        #
-        # Optimization: We derive total_num_tokens from the already all-reduced tokens_per_expert:
-        #   - tokens_per_expert.sum() = total valid token-expert assignments (global)
-        #   - Each valid token is assigned to exactly `topk` experts
-        #   - Therefore: total_valid_tokens = tokens_per_expert.sum() / topk
-        #
-        # This avoids an extra all_reduce operation.
-        if padding_mask is not None:
-            total_num_tokens = tokens_per_expert.sum() / self.topk
-        else:
-            total_num_tokens = num_tokens * self.tp_dp_cp_group.size()
 
         global_aux_loss = switch_load_balancing_loss_func(
             probs=scores_for_aux_loss,
@@ -481,7 +398,7 @@ def _apply_global_aux_loss(
             global_aux_loss,
             "global_load_balancing_loss",
             self.tp_dp_cp_group,
-            valid_token_count=num_tokens,
+            valid_token_count=local_num_tokens,
         )
         return probs
 
@@ -612,11 +529,8 @@ def _apply_expert_bias(
         if self.enable_expert_bias and torch.is_grad_enabled():
             with torch.no_grad():
                 if padding_mask is not None:
-                    # Only count valid tokens, exclude padding tokens
-                    # This logic is consistent with compute_tokens_per_expert
-                    mask_expanded = padding_mask.unsqueeze(-1)
-                    routing_map_masked = routing_map & mask_expanded
-                    self.local_tokens_per_expert += routing_map_masked.sum(dim=0)
+                    routing_map = routing_map & (~padding_mask)
+                    self.local_tokens_per_expert += routing_map.sum(dim=0)
                 else:
                     self.local_tokens_per_expert += routing_map.sum(dim=0)
 
@@ -639,12 +553,10 @@ def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = N
 
         # Flatten padding_mask to [num_tokens] if provided
         if padding_mask is not None:
-            padding_mask_flat = padding_mask.reshape(-1)
-        else:
-            padding_mask_flat = None
+            padding_mask = padding_mask.reshape(-1)
 
         # Apply Z-Loss
-        logits = self.apply_z_loss(logits, padding_mask=padding_mask_flat)
+        logits = self.apply_z_loss(logits, padding_mask=padding_mask)
 
         # Calculate probs and routing_map for token dispatching
         if self.routing_type == "sinkhorn":
@@ -681,10 +593,10 @@ def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = N
                 self.topk,
                 self.score_function,
                 fused=self.config.moe_router_fusion,
-                padding_mask=padding_mask_flat,
+                padding_mask=padding_mask,
             )
             probs = self._apply_aux_loss(
-                probs, scores_for_aux_loss, routing_map_for_aux_loss, padding_mask=padding_mask_flat
+                probs, scores_for_aux_loss, routing_map_for_aux_loss, padding_mask=padding_mask
             )
             probs = self._apply_seq_aux_loss(
                 probs,
@@ -692,14 +604,14 @@ def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = N
                 routing_map_for_aux_loss,
                 seq_length,
                 bsz,
-                padding_mask=padding_mask_flat,
+                padding_mask=padding_mask,
             )
             probs = self._apply_global_aux_loss(
-                probs, scores_for_aux_loss, routing_map_for_aux_loss, padding_mask=padding_mask_flat
+                probs, scores_for_aux_loss, routing_map_for_aux_loss, padding_mask=padding_mask
             )
 
         # Optionally apply expert bias
-        self._apply_expert_bias(routing_map, padding_mask=padding_mask_flat)
+        self._apply_expert_bias(routing_map, padding_mask=padding_mask)
 
         return probs, routing_map
 
diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py
index 93be17e9326..3e92694571a 100644
--- a/megatron/core/transformer/transformer_layer.py
+++ b/megatron/core/transformer/transformer_layer.py
@@ -597,19 +597,6 @@ def _forward_mlp(self, hidden_states, inference_context=None, padding_mask=None)
         """
         Perform a forward pass through the feed-forward layer.
 
-        For MoE layers, this method handles padding mask forwarding to ensure
-        padding tokens are excluded from auxiliary loss calculations.
-
-        Padding Mask Flow:
-            1. User provides padding_mask to GPTModel.forward() with shape [bsz, seq_length]
-            2. TransformerBlock passes it through to each TransformerLayer
-            3. This method passes it to MoELayer.forward() (for MoE layers only)
-            4. MoELayer uses PaddingMaskInfo to transform shape and handle masking
-
-        Padding Mask Semantics:
-            - True  = padding position (should be EXCLUDED from computation)
-            - False = valid position (should be INCLUDED in computation)
-
         Args:
             hidden_states (Tensor): Transformed hidden states before the MLP layernorm.
                 Shape [seq_length, batch_size, hidden_size].
@@ -653,24 +640,6 @@ def _forward_mlp(self, hidden_states, inference_context=None, padding_mask=None)
             and not isinstance(self.mlp, IdentityOp)
         )
 
-        # Helper function for MoE forward with padding_mask
-        def get_mlp_forward_fn():
-            """Return the appropriate MLP forward function based on layer type."""
-            if self.is_moe_layer and padding_mask is not None:
-
-                def mlp_forward(hidden_states):
-                    return self.mlp(hidden_states, padding_mask=padding_mask)
-
-                return mlp_forward
-            elif self.is_moe_layer:
-
-                def mlp_forward(hidden_states):
-                    return self.mlp(hidden_states)
-
-                return mlp_forward
-            else:
-                return self.mlp
-
         if (
             self.is_moe_layer
             and self.config.cuda_graph_impl == "transformer_engine"
@@ -684,21 +653,21 @@ def mlp_forward(hidden_states):
             cudagraph_outputs = self.mlp(pre_mlp_layernorm_output, padding_mask=padding_mask)
             return cudagraph_outputs + [residual]
         elif self.recompute_mlp:
-            mlp_fn = get_mlp_forward_fn()
             if self.config.fp8:
                 # import here to avoid circular import
                 from megatron.core.extensions.transformer_engine import te_checkpoint
 
                 mlp_output_with_bias = te_checkpoint(
-                    mlp_fn,
+                    self.mlp,
                     False,
                     tensor_parallel.random.get_cuda_rng_tracker,
                     self.pg_collection.tp,
                     pre_mlp_layernorm_output,
+                    padding_mask=padding_mask,
                 )
             else:
                 mlp_output_with_bias = tensor_parallel.checkpoint(
-                    mlp_fn, False, pre_mlp_layernorm_output
+                    self.mlp, False, pre_mlp_layernorm_output, padding_mask=padding_mask
                 )
         elif should_chunk_mlp_for_prefill:
             # Chunk input along sequence dimension
@@ -714,10 +683,7 @@ def mlp_forward(hidden_states):
             bias_output = torch.stack(bias_chunks, dim=0).sum(dim=0) if bias_chunks else None
             mlp_output_with_bias = (mlp_output, bias_output)
         else:
-            if self.is_moe_layer:
-                mlp_output_with_bias = self.mlp(pre_mlp_layernorm_output, padding_mask=padding_mask)
-            else:
-                mlp_output_with_bias = self.mlp(pre_mlp_layernorm_output)
+            mlp_output_with_bias = self.mlp(pre_mlp_layernorm_output, padding_mask=padding_mask)
 
         if self.recompute_pre_mlp_layernorm:
             # discard the output of the pre-mlp layernorm and register the recompute

From 9dd4373e35ddb0360ed215d28e89eb76b3531112 Mon Sep 17 00:00:00 2001
From: HaochenYuan <haocheny@nvidia.com>
Date: Thu, 11 Dec 2025 01:24:17 -0800
Subject: [PATCH 12/32] refactor

---
 megatron/core/transformer/moe/moe_layer.py | 11 -----------
 megatron/core/transformer/moe/moe_utils.py |  3 ---
 megatron/core/transformer/moe/router.py    |  4 +---
 3 files changed, 1 insertion(+), 17 deletions(-)

diff --git a/megatron/core/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py
index ca7ead48540..098e17f765c 100644
--- a/megatron/core/transformer/moe/moe_layer.py
+++ b/megatron/core/transformer/moe/moe_layer.py
@@ -12,7 +12,6 @@
 from megatron.core.transformer.moe.moe_utils import (
     MoECudaGraphPartialCaptureSignal,
     MoECudaGraphTensorStore,
-    PaddingMaskInfo,
     get_default_pg_collection,
     maybe_skip_or_early_return_by_cudagraph,
 )
@@ -184,16 +183,6 @@ def route(self, hidden_states: torch.Tensor, padding_mask: Optional[torch.Tensor
 
         This method uses the router to determine which experts to send each token to,
         producing routing probabilities and a mapping.
-
-        Note: The padding_mask transformation from [bsz, seq_length] to [seq_length, bsz]
-        is done in the forward() method before calling this method.
-
-        Args:
-            hidden_states (torch.Tensor): Input tensor with shape [seq_length, bsz, hidden_size].
-            padding_mask (torch.Tensor, optional): Boolean mask indicating padding positions.
-                Shape = [seq_length, bsz] (after transformation in forward()).
-                True = padding (exclude), False = valid (include).
-                Defaults to None (all tokens are valid).
         """
         probs, routing_map = self.router(hidden_states, padding_mask=padding_mask)
         return probs, routing_map
diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py
index 6c1342bda80..47c14727fa4 100644
--- a/megatron/core/transformer/moe/moe_utils.py
+++ b/megatron/core/transformer/moe/moe_utils.py
@@ -91,9 +91,6 @@ def switch_load_balancing_loss_func(
         topk (int): The number of experts selected for each token.
         num_experts (int): The number of experts.
         moe_aux_loss_coeff (float): The coefficient for the auxiliary loss.
-        padding_mask (torch.Tensor, optional): Boolean mask indicating padding positions.
-                                               Shape [num_tokens]. True = padding (exclude),
-                                               False = valid (include). Defaults to None.
     Returns:
         torch.Tensor: The auxiliary loss for load balancing.
     """
diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py
index 7dfd780a164..e891837dc84 100644
--- a/megatron/core/transformer/moe/router.py
+++ b/megatron/core/transformer/moe/router.py
@@ -530,9 +530,7 @@ def _apply_expert_bias(
             with torch.no_grad():
                 if padding_mask is not None:
                     routing_map = routing_map & (~padding_mask)
-                    self.local_tokens_per_expert += routing_map.sum(dim=0)
-                else:
-                    self.local_tokens_per_expert += routing_map.sum(dim=0)
+                self.local_tokens_per_expert += routing_map.sum(dim=0)
 
     def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = None):
         """Top-k routing function

From 2a9f3c5633b038e6de9e49d9e70079420f139bff Mon Sep 17 00:00:00 2001
From: HaochenYuan <haocheny@nvidia.com>
Date: Thu, 11 Dec 2025 01:25:16 -0800
Subject: [PATCH 13/32] refactor

---
 megatron/core/transformer/moe/moe_utils.py | 10 +++---
 megatron/core/transformer/moe/router.py    | 38 ++++++++++++----------
 2 files changed, 27 insertions(+), 21 deletions(-)

diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py
index 47c14727fa4..aa7b5da21a8 100644
--- a/megatron/core/transformer/moe/moe_utils.py
+++ b/megatron/core/transformer/moe/moe_utils.py
@@ -2,11 +2,11 @@
 from dataclasses import dataclass
 from typing import List, Optional, Union
 
-from megatron.core.tensor_parallel.mappings import reduce_from_tensor_model_parallel_region
 import torch
 
 from megatron.core import parallel_state
 from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.tensor_parallel.mappings import reduce_from_tensor_model_parallel_region
 from megatron.core.transformer.cuda_graphs import is_graph_capturing
 
 try:
@@ -179,11 +179,12 @@ def get_capacity(num_tokens: int, num_experts: int, capacity_factor: float, min_
         capacity = min_capacity
     return capacity
 
+
 def get_tokens_per_expert_and_token_count(
-    routing_map: torch.Tensor, 
-    reduce_group: torch.distributed.ProcessGroup, 
+    routing_map: torch.Tensor,
+    reduce_group: torch.distributed.ProcessGroup,
     topk: int = None,
-    with_padding_mask: bool = False
+    with_padding_mask: bool = False,
 ) -> torch.Tensor:
     local_tokens_per_expert = routing_map.sum(dim=0)
     global_tokens_per_expert = reduce_from_tensor_model_parallel_region(
@@ -197,6 +198,7 @@ def get_tokens_per_expert_and_token_count(
         total_num_tokens = local_num_tokens * reduce_group.size()
     return global_tokens_per_expert, local_num_tokens, total_num_tokens
 
+
 class MoEAuxLossAutoScaler(torch.autograd.Function):
     """An AutoScaler that triggers the backward pass and scales the grad for auxiliary loss."""
 
diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py
index e891837dc84..0184e8a569a 100644
--- a/megatron/core/transformer/moe/router.py
+++ b/megatron/core/transformer/moe/router.py
@@ -280,11 +280,13 @@ def _apply_aux_loss(
         if aux_loss_coeff == 0:
             return probs
 
-        global_tokens_per_expert, local_num_tokens, total_num_tokens = get_tokens_per_expert_and_token_count(
-            routing_map=routing_map,
-            reduce_group=self.tp_cp_group,
-            topk=self.topk,
-            with_padding_mask=with_padding_mask,
+        global_tokens_per_expert, local_num_tokens, total_num_tokens = (
+            get_tokens_per_expert_and_token_count(
+                routing_map=routing_map,
+                reduce_group=self.tp_cp_group,
+                topk=self.topk,
+                with_padding_mask=with_padding_mask,
+            )
         )
         aux_loss = switch_load_balancing_loss_func(
             probs=scores_for_aux_loss,
@@ -327,11 +329,13 @@ def _apply_seq_aux_loss(
 
         scores_for_aux_loss = scores_for_aux_loss.reshape(seq_length, -1)
 
-        global_tokens_per_expert, local_num_tokens, total_num_tokens = get_tokens_per_expert_and_token_count(
-            routing_map=routing_map,
-            reduce_group=self.tp_cp_group,
-            with_padding_mask=with_padding_mask,
-            topk=self.topk,
+        global_tokens_per_expert, local_num_tokens, total_num_tokens = (
+            get_tokens_per_expert_and_token_count(
+                routing_map=routing_map,
+                reduce_group=self.tp_cp_group,
+                with_padding_mask=with_padding_mask,
+                topk=self.topk,
+            )
         )
 
         aux_loss = (
@@ -347,7 +351,6 @@ def _apply_seq_aux_loss(
             / bsz
         )
 
-
         probs = self.attach_and_log_load_balancing_loss(
             probs,
             seq_aux_loss_coeff,
@@ -371,17 +374,18 @@ def _apply_global_aux_loss(
             return probs
 
         # Use unified function to compute tokens_per_expert and num_tokens
-        global_tokens_per_expert, local_num_tokens, total_num_tokens = get_tokens_per_expert_and_token_count(
-            routing_map=routing_map,
-            reduce_group=self.tp_dp_cp_group,
-            with_padding_mask=with_padding_mask,
-            topk=self.topk,
+        global_tokens_per_expert, local_num_tokens, total_num_tokens = (
+            get_tokens_per_expert_and_token_count(
+                routing_map=routing_map,
+                reduce_group=self.tp_dp_cp_group,
+                with_padding_mask=with_padding_mask,
+                topk=self.topk,
+            )
         )
         self.global_tokens_per_expert += global_tokens_per_expert
         self.ga_steps += 1
         averated_tokens_per_expert = self.global_tokens_per_expert / self.ga_steps
 
-
         global_aux_loss = switch_load_balancing_loss_func(
             probs=scores_for_aux_loss,
             tokens_per_expert=averated_tokens_per_expert,

From cc34bcd26a7e4c51e7a420f6c657999b690b740c Mon Sep 17 00:00:00 2001
From: HaochenYuan <haocheny@nvidia.com>
Date: Thu, 11 Dec 2025 01:28:15 -0800
Subject: [PATCH 14/32] refactor

---
 megatron/core/transformer/moe/router.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py
index 0184e8a569a..8134a3a91c7 100644
--- a/megatron/core/transformer/moe/router.py
+++ b/megatron/core/transformer/moe/router.py
@@ -6,7 +6,6 @@
 import torch
 
 from megatron.core.jit import jit_fuser
-from megatron.core.tensor_parallel import reduce_from_tensor_model_parallel_region
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.transformer.moe.moe_utils import (
     MoEAuxLossAutoScaler,

From 886a613439297e7e11ab3f82b1d2bfc22d5c058c Mon Sep 17 00:00:00 2001
From: HaochenYuan <haocheny@nvidia.com>
Date: Thu, 11 Dec 2025 04:46:01 -0800
Subject: [PATCH 15/32] fix bug

---
 megatron/core/transformer/moe/moe_utils.py |  3 +++
 megatron/core/transformer/moe/router.py    | 18 +++++++++++++-----
 2 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py
index aa7b5da21a8..23a97d7ff84 100644
--- a/megatron/core/transformer/moe/moe_utils.py
+++ b/megatron/core/transformer/moe/moe_utils.py
@@ -186,6 +186,9 @@ def get_tokens_per_expert_and_token_count(
     topk: int = None,
     with_padding_mask: bool = False,
 ) -> torch.Tensor:
+    """
+    Compute global_tokens_per_expert, local_num_tokens and total_num_tokens with optional padding mask.
+    """
     local_tokens_per_expert = routing_map.sum(dim=0)
     global_tokens_per_expert = reduce_from_tensor_model_parallel_region(
         local_tokens_per_expert, reduce_group
diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py
index 8134a3a91c7..3f439046014 100644
--- a/megatron/core/transformer/moe/router.py
+++ b/megatron/core/transformer/moe/router.py
@@ -327,6 +327,7 @@ def _apply_seq_aux_loss(
             return probs
 
         scores_for_aux_loss = scores_for_aux_loss.reshape(seq_length, -1)
+        routing_map = routing_map.reshape(seq_length, -1)
 
         global_tokens_per_expert, local_num_tokens, total_num_tokens = (
             get_tokens_per_expert_and_token_count(
@@ -336,7 +337,9 @@ def _apply_seq_aux_loss(
                 topk=self.topk,
             )
         )
-
+        if with_padding_mask:
+            local_num_tokens //= bsz
+            total_num_tokens //= bsz
         aux_loss = (
             switch_load_balancing_loss_func(
                 probs=scores_for_aux_loss,
@@ -393,7 +396,6 @@ def _apply_global_aux_loss(
             num_experts=self.config.num_moe_experts,
             moe_aux_loss_coeff=global_aux_loss_coeff,
             fused=self.config.moe_router_fusion,
-            padding_mask=None,
         )
         probs = self.attach_and_log_load_balancing_loss(
             probs,
@@ -597,7 +599,10 @@ def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = N
                 padding_mask=padding_mask,
             )
             probs = self._apply_aux_loss(
-                probs, scores_for_aux_loss, routing_map_for_aux_loss, padding_mask=padding_mask
+                probs,
+                scores_for_aux_loss,
+                routing_map_for_aux_loss,
+                with_padding_mask=padding_mask is not None,
             )
             probs = self._apply_seq_aux_loss(
                 probs,
@@ -605,10 +610,13 @@ def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = N
                 routing_map_for_aux_loss,
                 seq_length,
                 bsz,
-                padding_mask=padding_mask,
+                with_padding_mask=padding_mask is not None,
             )
             probs = self._apply_global_aux_loss(
-                probs, scores_for_aux_loss, routing_map_for_aux_loss, padding_mask=padding_mask
+                probs,
+                scores_for_aux_loss,
+                routing_map_for_aux_loss,
+                with_padding_mask=padding_mask is not None,
             )
 
         # Optionally apply expert bias

From 97de0670e6c40e38b89f850e13672828e1c8cd9b Mon Sep 17 00:00:00 2001
From: HaochenYuan <haocheny@nvidia.com>
Date: Thu, 11 Dec 2025 04:56:18 -0800
Subject: [PATCH 16/32] fix bug in seq_aux_loss

---
 megatron/core/transformer/moe/router.py | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py
index 3f439046014..416c882ff97 100644
--- a/megatron/core/transformer/moe/router.py
+++ b/megatron/core/transformer/moe/router.py
@@ -334,12 +334,10 @@ def _apply_seq_aux_loss(
                 routing_map=routing_map,
                 reduce_group=self.tp_cp_group,
                 with_padding_mask=with_padding_mask,
-                topk=self.topk,
+                topk=self.topk*bsz,
             )
         )
-        if with_padding_mask:
-            local_num_tokens //= bsz
-            total_num_tokens //= bsz
+
         aux_loss = (
             switch_load_balancing_loss_func(
                 probs=scores_for_aux_loss,

From fdb731173bc1faa80b1aa6370af06483d616aef3 Mon Sep 17 00:00:00 2001
From: HaochenYuan <haocheny@nvidia.com>
Date: Thu, 11 Dec 2025 04:59:30 -0800
Subject: [PATCH 17/32] fix bug in seq_aux_loss

---
 megatron/core/transformer/moe/router.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py
index 416c882ff97..1c502e212ad 100644
--- a/megatron/core/transformer/moe/router.py
+++ b/megatron/core/transformer/moe/router.py
@@ -334,7 +334,7 @@ def _apply_seq_aux_loss(
                 routing_map=routing_map,
                 reduce_group=self.tp_cp_group,
                 with_padding_mask=with_padding_mask,
-                topk=self.topk*bsz,
+                topk=self.topk * bsz,
             )
         )
 

From c30213212a25e9c60d32d09cb724872a157821c6 Mon Sep 17 00:00:00 2001
From: HaochenYuan <haocheny@nvidia.com>
Date: Thu, 11 Dec 2025 05:04:31 -0800
Subject: [PATCH 18/32] fix linting

---
 megatron/core/transformer/moe/moe_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py
index 23a97d7ff84..8f904207ad4 100644
--- a/megatron/core/transformer/moe/moe_utils.py
+++ b/megatron/core/transformer/moe/moe_utils.py
@@ -187,7 +187,7 @@ def get_tokens_per_expert_and_token_count(
     with_padding_mask: bool = False,
 ) -> torch.Tensor:
     """
-    Compute global_tokens_per_expert, local_num_tokens and total_num_tokens with optional padding mask.
+    Compute global_tokens_per_expert, local_num_tokens and total_num_tokens with padding mask.
     """
     local_tokens_per_expert = routing_map.sum(dim=0)
     global_tokens_per_expert = reduce_from_tensor_model_parallel_region(

From aa3ac99cee4537cfa2aeed54f5f0b9be8a9915f6 Mon Sep 17 00:00:00 2001
From: HaochenYuan <haocheny@nvidia.com>
Date: Thu, 11 Dec 2025 05:08:08 -0800
Subject: [PATCH 19/32] add Copyright header

---
 megatron/core/transformer/moe/moe_utils.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py
index 8f904207ad4..88323501853 100644
--- a/megatron/core/transformer/moe/moe_utils.py
+++ b/megatron/core/transformer/moe/moe_utils.py
@@ -1,3 +1,4 @@
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 import math
 from dataclasses import dataclass
 from typing import List, Optional, Union

From 39e3db4f77a3cb45e64bf5cb15e73e1766a99f9d Mon Sep 17 00:00:00 2001
From: HaochenYuan <haocheny@nvidia.com>
Date: Thu, 11 Dec 2025 05:10:16 -0800
Subject: [PATCH 20/32] add Copyright header

---
 megatron/core/transformer/transformer_layer.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py
index 3e92694571a..d270b87a532 100644
--- a/megatron/core/transformer/transformer_layer.py
+++ b/megatron/core/transformer/transformer_layer.py
@@ -651,6 +651,7 @@ def _forward_mlp(self, hidden_states, inference_context=None, padding_mask=None)
                 not self.recompute_pre_mlp_layernorm
             ), "Recomputation is not supported for CUDA graph."
             cudagraph_outputs = self.mlp(pre_mlp_layernorm_output, padding_mask=padding_mask)
+            nvtx_range_pop(suffix="mlp")
             return cudagraph_outputs + [residual]
         elif self.recompute_mlp:
             if self.config.fp8:

From cebf69df870f6f3a01784ccb21be31fcb45e3fdd Mon Sep 17 00:00:00 2001
From: HaochenYuan <haocheny@nvidia.com>
Date: Thu, 11 Dec 2025 22:45:21 -0800
Subject: [PATCH 21/32] add padding_mask in fusedMLP

---
 megatron/core/extensions/transformer_engine.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py
index e807ee54fbf..b1dcd1a24f9 100644
--- a/megatron/core/extensions/transformer_engine.py
+++ b/megatron/core/extensions/transformer_engine.py
@@ -1762,7 +1762,7 @@ def forward_post_hook(module, *_) -> None:
                     "TEFusedMLP module does not support submodules with post-backward hooks"
                 )
 
-        def forward(self, hidden_states: torch.Tensor) -> Tuple[Tensor, Optional[Tensor]]:
+        def forward(self, hidden_states: torch.Tensor, **kwargs) -> Tuple[Tensor, Optional[Tensor]]:
             """Forward."""
 
             # Construct fused impl if needed

From bc1a37c6da40b5efe32d608b5382be6ef3ea0b44 Mon Sep 17 00:00:00 2001
From: HaochenYuan <haocheny@nvidia.com>
Date: Thu, 18 Dec 2025 04:21:36 -0800
Subject: [PATCH 22/32] slice padding_mask in SP

---
 megatron/core/models/gpt/gpt_model.py | 15 +++++++++++----
 1 file changed, 11 insertions(+), 4 deletions(-)

diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py
index ad4fb864333..d0633ba320b 100644
--- a/megatron/core/models/gpt/gpt_model.py
+++ b/megatron/core/models/gpt/gpt_model.py
@@ -284,6 +284,7 @@ def _preprocess(
         decoder_input: Tensor = None,
         inference_context: BaseInferenceContext = None,
         packed_seq_params: PackedSeqParams = None,
+        padding_mask: Optional[Tensor] = None,
     ):
         """Preprocesses inputs for the transformer decoder.
 
@@ -301,6 +302,11 @@ def _preprocess(
             pass
         elif self.pre_process:
             decoder_input = self.embedding(input_ids=input_ids, position_ids=position_ids)
+            if padding_mask is not None and self.config.sequence_parallel:
+                # padding_mask: [bsz, seq] -> [seq, bsz] -> scatter -> [seq/tp, bsz] -> [bsz, seq/tp]
+                padding_mask = tensor_parallel.scatter_to_sequence_parallel_region(
+                    padding_mask.transpose(0, 1).contiguous()
+                ).transpose(0, 1).contiguous()
         else:
             # intermediate stage of pipeline
             # decoder will get hidden_states from encoder.input_tensor
@@ -403,6 +409,7 @@ def _preprocess(
             rotary_pos_cos,
             rotary_pos_sin,
             sequence_len_offset,
+            padding_mask,
         )
         if rotary_pos_cos_sin is not None:
             # only in the case of flashinfer fused rope will we
@@ -472,13 +479,13 @@ def forward(
             decoder_input=decoder_input,
             inference_context=inference_context,
             packed_seq_params=packed_seq_params,
+            padding_mask=padding_mask,
         )
 
-        (decoder_input, rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, sequence_len_offset) = (
-            preproc_output[:5]
-        )
+        (decoder_input, rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, sequence_len_offset,
+         padding_mask) = preproc_output[:6]
 
-        rotary_pos_cos_sin = preproc_output[5] if len(preproc_output) == 6 else None
+        rotary_pos_cos_sin = preproc_output[6] if len(preproc_output) == 7 else None
 
         # Run decoder.
         hidden_states = self.decoder(

From bc78b601776747c4ce517673d5c81e47b276d71a Mon Sep 17 00:00:00 2001
From: HaochenYuan <haocheny@nvidia.com>
Date: Thu, 18 Dec 2025 04:26:54 -0800
Subject: [PATCH 23/32] fix linting

---
 megatron/core/models/gpt/gpt_model.py | 21 +++++++++++++++------
 1 file changed, 15 insertions(+), 6 deletions(-)

diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py
index d0633ba320b..b5f11441335 100644
--- a/megatron/core/models/gpt/gpt_model.py
+++ b/megatron/core/models/gpt/gpt_model.py
@@ -303,10 +303,13 @@ def _preprocess(
         elif self.pre_process:
             decoder_input = self.embedding(input_ids=input_ids, position_ids=position_ids)
             if padding_mask is not None and self.config.sequence_parallel:
-                # padding_mask: [bsz, seq] -> [seq, bsz] -> scatter -> [seq/tp, bsz] -> [bsz, seq/tp]
-                padding_mask = tensor_parallel.scatter_to_sequence_parallel_region(
-                    padding_mask.transpose(0, 1).contiguous()
-                ).transpose(0, 1).contiguous()
+                padding_mask = (
+                    tensor_parallel.scatter_to_sequence_parallel_region(
+                        padding_mask.transpose(0, 1).contiguous()
+                    )
+                    .transpose(0, 1)
+                    .contiguous()
+                )
         else:
             # intermediate stage of pipeline
             # decoder will get hidden_states from encoder.input_tensor
@@ -482,8 +485,14 @@ def forward(
             padding_mask=padding_mask,
         )
 
-        (decoder_input, rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, sequence_len_offset,
-         padding_mask) = preproc_output[:6]
+        (
+            decoder_input,
+            rotary_pos_emb,
+            rotary_pos_cos,
+            rotary_pos_sin,
+            sequence_len_offset,
+            padding_mask,
+        ) = preproc_output[:6]
 
         rotary_pos_cos_sin = preproc_output[6] if len(preproc_output) == 7 else None
 

From a9f02b1eb1f5049e399b48ca16d004fb97307d47 Mon Sep 17 00:00:00 2001
From: HaochenYuan <haocheny@nvidia.com>
Date: Thu, 18 Dec 2025 05:14:52 -0800
Subject: [PATCH 24/32] modify preprocess in 1f1b

---
 .../core/models/gpt/fine_grained_callables.py | 21 ++++++++++++-------
 1 file changed, 14 insertions(+), 7 deletions(-)

diff --git a/megatron/core/models/gpt/fine_grained_callables.py b/megatron/core/models/gpt/fine_grained_callables.py
index 60094976a9a..ceb11dbcc11 100644
--- a/megatron/core/models/gpt/fine_grained_callables.py
+++ b/megatron/core/models/gpt/fine_grained_callables.py
@@ -120,13 +120,19 @@ def forward_impl(self):
         if not self.gpt_model.pre_process:
             self.chunk_state.decoder_input = self.gpt_model.decoder.input_tensor
         # Run GPTModel._preprocess
-        decoder_input, rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, sequence_len_offset = (
-            self.gpt_model._preprocess(
-                input_ids=self.chunk_state.input_ids,
-                position_ids=self.chunk_state.position_ids,
-                decoder_input=self.chunk_state.decoder_input,
-                packed_seq_params=self.chunk_state.packed_seq_params,
-            )
+        (
+            decoder_input,
+            rotary_pos_emb,
+            rotary_pos_cos,
+            rotary_pos_sin,
+            sequence_len_offset,
+            padding_mask,
+        ) = self.gpt_model._preprocess(
+            input_ids=self.chunk_state.input_ids,
+            position_ids=self.chunk_state.position_ids,
+            decoder_input=self.chunk_state.decoder_input,
+            packed_seq_params=self.chunk_state.packed_seq_params,
+            padding_mask=self.chunk_state.padding_mask,
         )
 
         # Saved for later use
@@ -135,6 +141,7 @@ def forward_impl(self):
         self.chunk_state.rotary_pos_cos = rotary_pos_cos
         self.chunk_state.rotary_pos_sin = rotary_pos_sin
         self.chunk_state.sequence_len_offset = sequence_len_offset
+        self.chunk_state.padding_mask = padding_mask
         return decoder_input
 
 

From 9ae9fc98c4846b9a55aff99d6e4282b1bd7736ce Mon Sep 17 00:00:00 2001
From: HaochenYuan <haocheny@nvidia.com>
Date: Thu, 18 Dec 2025 06:10:04 -0800
Subject: [PATCH 25/32] modify preprocess in 1f1b&mtp

---
 megatron/core/models/common/model_chunk_schedule_plan.py | 2 ++
 megatron/core/transformer/moe/moe_utils.py               | 1 +
 tests/unit_tests/a2a_overlap/test_schedule_layer_1f1b.py | 2 +-
 3 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/megatron/core/models/common/model_chunk_schedule_plan.py b/megatron/core/models/common/model_chunk_schedule_plan.py
index 486a498dd73..07bab1cb486 100644
--- a/megatron/core/models/common/model_chunk_schedule_plan.py
+++ b/megatron/core/models/common/model_chunk_schedule_plan.py
@@ -305,6 +305,7 @@ def __init__(
         extra_block_kwargs=None,
         runtime_gather_output: Optional[bool] = None,
         loss_mask: Optional[Tensor] = None,
+        padding_mask=None,
     ):
         """Initialize the schedule plan of all Transformer layers' sub-modules.
 
@@ -347,6 +348,7 @@ def __init__(
         self._model_chunk_state.mtp_hidden_states = None
         self._model_chunk_state.loss_mask = loss_mask
         self._model_chunk_state.packed_seq_params = packed_seq_params
+        self._model_chunk_state.padding_mask = padding_mask
         self._model_chunk_state.extra_block_kwargs = extra_block_kwargs
         self._model_chunk_state.runtime_gather_output = runtime_gather_output
         self._model_chunk_state.model = model
diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py
index 4c23e009437..5937a569784 100644
--- a/megatron/core/transformer/moe/moe_utils.py
+++ b/megatron/core/transformer/moe/moe_utils.py
@@ -9,6 +9,7 @@
 from megatron.core.fp4_utils import get_fp4_align_size
 from megatron.core.fp8_utils import get_fp8_align_size
 from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.tensor_parallel import get_cuda_rng_tracker, get_expert_parallel_rng_tracker_name
 from megatron.core.tensor_parallel.mappings import reduce_from_tensor_model_parallel_region
 from megatron.core.transformer.cuda_graphs import is_graph_capturing
 from megatron.core.transformer.enums import CudaGraphScope
diff --git a/tests/unit_tests/a2a_overlap/test_schedule_layer_1f1b.py b/tests/unit_tests/a2a_overlap/test_schedule_layer_1f1b.py
index 7fb97f6e586..b3f4c56409f 100644
--- a/tests/unit_tests/a2a_overlap/test_schedule_layer_1f1b.py
+++ b/tests/unit_tests/a2a_overlap/test_schedule_layer_1f1b.py
@@ -502,7 +502,7 @@ def test_mtp_layer_overlap(self, dispatcher_type, fp8_flag):
             position_ids = torch.tensor(data, dtype=torch.int64).repeat((1, 1)).cuda()
             attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool).cuda()
             # get rotary pos emb
-            _, rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, _ = gpt_model._preprocess(
+            _, rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, _, _padding_mask = gpt_model._preprocess(
                 input_ids, position_ids
             )
             # reset model

From c1195c39275d22557f11925cb552775972cca046 Mon Sep 17 00:00:00 2001
From: HaochenYuan <haocheny@nvidia.com>
Date: Thu, 18 Dec 2025 06:13:45 -0800
Subject: [PATCH 26/32] fix linting

---
 tests/unit_tests/a2a_overlap/test_schedule_layer_1f1b.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/unit_tests/a2a_overlap/test_schedule_layer_1f1b.py b/tests/unit_tests/a2a_overlap/test_schedule_layer_1f1b.py
index b3f4c56409f..5ec096e5a04 100644
--- a/tests/unit_tests/a2a_overlap/test_schedule_layer_1f1b.py
+++ b/tests/unit_tests/a2a_overlap/test_schedule_layer_1f1b.py
@@ -502,8 +502,8 @@ def test_mtp_layer_overlap(self, dispatcher_type, fp8_flag):
             position_ids = torch.tensor(data, dtype=torch.int64).repeat((1, 1)).cuda()
             attention_mask = torch.ones((1, 1, seq_len, seq_len), dtype=bool).cuda()
             # get rotary pos emb
-            _, rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, _, _padding_mask = gpt_model._preprocess(
-                input_ids, position_ids
+            _, rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, _, _padding_mask = (
+                gpt_model._preprocess(input_ids, position_ids)
             )
             # reset model
             params = reset_model(gpt_model)

From cae40958ae0c0a3a9574a205710da3f53d91eb2f Mon Sep 17 00:00:00 2001
From: HaochenYuan <haocheny@nvidia.com>
Date: Tue, 23 Dec 2025 01:45:22 -0800
Subject: [PATCH 27/32] add UT in 1f1b & SP-S

---
 megatron/core/models/gpt/gpt_model.py         |   8 ++
 .../a2a_overlap/test_schedule_chunk_1f1b.py   | 118 +++++++++++++++++-
 .../transformer/moe/test_aux_loss.py          |  11 +-
 3 files changed, 132 insertions(+), 5 deletions(-)

diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py
index b5f11441335..9e70c677226 100644
--- a/megatron/core/models/gpt/gpt_model.py
+++ b/megatron/core/models/gpt/gpt_model.py
@@ -301,6 +301,11 @@ def _preprocess(
         if decoder_input is not None:
             pass
         elif self.pre_process:
+            if padding_mask is not None:
+                assert padding_mask.shape == input_ids.shape, (
+                    f"padding_mask shape {padding_mask.shape} does not match "
+                    f"input_ids shape {input_ids.shape}"
+                )
             decoder_input = self.embedding(input_ids=input_ids, position_ids=position_ids)
             if padding_mask is not None and self.config.sequence_parallel:
                 padding_mask = (
@@ -745,6 +750,7 @@ def build_schedule_plan(
         runtime_gather_output: Optional[bool] = None,
         inference_params: Optional[BaseInferenceContext] = None,
         loss_mask: Optional[Tensor] = None,
+        padding_mask: Optional[Tensor] = None,
     ):
         """Builds a computation schedule plan for the model.
 
@@ -770,6 +776,7 @@ def build_schedule_plan(
             inference_params (InferenceParams, optional):
                 Parameters for inference. Defaults to None.
             loss_mask (Optional[Tensor], optional): Loss mask. Defaults to None.
+            padding_mask (Optional[Tensor], optional): Padding mask. Defaults to None.
 
         Returns:
             TransformerModelChunkSchedulePlan: The model chunk schedule plan.
@@ -791,6 +798,7 @@ def build_schedule_plan(
             extra_block_kwargs,
             runtime_gather_output,
             loss_mask,
+            padding_mask,
         )
 
     def sharded_state_dict(
diff --git a/tests/unit_tests/a2a_overlap/test_schedule_chunk_1f1b.py b/tests/unit_tests/a2a_overlap/test_schedule_chunk_1f1b.py
index 81e61a3404a..ad27187b63e 100644
--- a/tests/unit_tests/a2a_overlap/test_schedule_chunk_1f1b.py
+++ b/tests/unit_tests/a2a_overlap/test_schedule_chunk_1f1b.py
@@ -23,7 +23,7 @@
 from tests.unit_tests.test_utilities import Utils
 
 
-def build_model(config):
+def build_model(config, use_padding_mask=False):
     seq_len = 32
     max_seq_len = 300
     # ids = random.sample([i for i in range(max_seq_len)], seq_len)
@@ -39,6 +39,12 @@ def build_model(config):
         "attention_mask": torch.ones((1, 1, seq_len, seq_len), dtype=bool).cuda(),
     }
 
+    # Optionally add padding_mask with same shape as input_ids
+    if use_padding_mask:
+        padding_mask = torch.zeros((1, seq_len), dtype=torch.bool).cuda()
+        padding_mask[0, -8:] = True
+        data["padding_mask"] = padding_mask
+
     # build layer spec
     transformer_layer_spec = get_gpt_decoder_block_spec(config=config, use_transformer_engine=True)
     mtp_block_spec = get_gpt_mtp_block_spec(config, transformer_layer_spec.layer_specs[-1], True)
@@ -48,7 +54,7 @@ def build_model(config):
         config=config,
         transformer_layer_spec=transformer_layer_spec,
         mtp_block_spec=mtp_block_spec,
-        vocab_size=100,
+        vocab_size=128,
         pre_process=True,
         post_process=True,
         max_sequence_length=max_seq_len,
@@ -174,3 +180,111 @@ def test_1f1b_schedule_model_chunk(self, mtp_layers, dispatcher_type, fp8_flag,
                 gpt_models[i] = None
             gc.collect()
             torch.cuda.empty_cache()
+
+    @pytest.mark.skipif(not is_te_min_version("1.9.0.dev0"), reason="Requires TE >= 1.9.0.dev0")
+    @pytest.mark.parametrize("dispatcher_type", get_valid_token_dispatcher_types())
+    @pytest.mark.parametrize("layers", [[2, 1], [1, 1]])
+    @pytest.mark.parametrize("tp_size", [1, 2, 4, 8])
+    def test_1f1b_schedule_model_chunk_with_padding_mask(
+        self, dispatcher_type, layers, parallel_config
+    ):
+        """
+        Verifies all-to-all overlap optimization with padding_mask produces
+        the same results as the reference implementation with various TP/EP/CP combinations.
+        """
+        # Re-initialize model parallel with the specified configuration
+        Utils.destroy_model_parallel()
+        Utils.initialize_model_parallel(
+            tensor_model_parallel_size=parallel_config["tp"],
+            pipeline_model_parallel_size=1,
+            expert_model_parallel_size=4,
+            expert_tensor_parallel_size=1,
+        )
+        set_streams()
+
+        microbatches = 1
+
+        gpt_models = []
+        schedule_plans = []
+        ref_captures = []
+        datas = []
+
+        # create TransformerConfig
+        extra_kwargs = {
+            "moe_token_dispatcher_type": dispatcher_type,
+            "tensor_model_parallel_size": parallel_config["tp"],
+            "sequence_parallel": parallel_config["tp"] > 1,
+        }
+        if dispatcher_type == "flex":
+            extra_kwargs["moe_flex_dispatcher_backend"] = "deepep"
+            extra_kwargs["moe_router_dtype"] = "fp32"
+        with deterministic_mode():
+            for layer_num in layers:
+                output_tensors = []
+                # build config
+                config = get_test_config(num_layers=layer_num, extra_kwargs=extra_kwargs)
+                # build model with padding_mask
+                gpt_model, schedule_plan, data = build_model(config, use_padding_mask=True)
+                gpt_model.cuda()
+                gpt_models.append(gpt_model)
+                datas.append(data)
+                schedule_plans.append(schedule_plan)
+
+                # run reference
+                for _ in range(microbatches):
+                    loss = gpt_model.forward(**data)
+                    loss = float16_to_fp32(loss)
+                    loss.backward(torch.ones_like(loss))
+                    output_tensors.append(loss)
+
+                capture = {"outputs": output_tensors}
+                for name, param in gpt_model.named_parameters():
+                    capture[name] = param.grad
+                ref_captures.append(capture)
+                gpt_model.zero_grad()
+            assert gpt_models[0].embedding is not None
+            assert gpt_models[1].embedding is not None
+            # run a2a overlap
+            capture_0 = {"outputs": []}
+            capture_1 = {"outputs": []}
+            a2a_captures = [capture_0, capture_1]
+            for i in range(microbatches):
+                # 1st forward
+                if i > 0:
+                    assert (
+                        schedule_plans[0].pre_process is None
+                    ), "pre_process should be released after backward"
+                    schedule_plans[0] = gpt_models[0].build_schedule_plan(**datas[0])
+                    schedule_plans[1] = gpt_models[1].build_schedule_plan(**datas[1])
+                f_input_0 = TransformerModelChunkSchedulePlan.run(schedule_plans[0], None)
+                capture_0["outputs"].append(f_input_0)
+                # overlap
+                f_input_1 = TransformerModelChunkSchedulePlan.run(
+                    schedule_plans[1], schedule_plans[0], b_grad=torch.ones_like(f_input_0)
+                )
+                capture_1["outputs"].append(f_input_1)
+                # last backward
+                TransformerModelChunkSchedulePlan.run(
+                    None, schedule_plans[1], b_grad=torch.ones_like(f_input_1)
+                )
+            for i in range(len(gpt_models)):
+                for name, param in gpt_models[i].named_parameters():
+                    a2a_captures[i][name] = param.grad
+
+            # compare results
+            for i in range(len(ref_captures)):
+                comp_res = compare_captures(ref_captures[i], a2a_captures[i], True, True)
+                assert comp_res[0], f"[rank {torch.distributed.get_rank()}] {comp_res[1]}"
+
+            # release resources is necessary, otherwise later testcases will oom
+            for i in range(len(schedule_plans)):
+                schedule_plans[i] = None
+                ref_captures[i] = None
+                a2a_captures[i] = None
+                for k in datas[i]:
+                    datas[i][k] = None
+                datas[i] = None
+                gpt_models[i].zero_grad()
+                gpt_models[i] = None
+            gc.collect()
+            torch.cuda.empty_cache()
diff --git a/tests/unit_tests/transformer/moe/test_aux_loss.py b/tests/unit_tests/transformer/moe/test_aux_loss.py
index cee2162ade2..319fbbea5c3 100644
--- a/tests/unit_tests/transformer/moe/test_aux_loss.py
+++ b/tests/unit_tests/transformer/moe/test_aux_loss.py
@@ -581,7 +581,7 @@ def test_force_balanced_aux_loss(self, tp_size, ep_size, cp_size):
 class TestPaddingMaskAuxLoss:
     """Test padding mask support in various aux loss types."""
 
-    def setup_model_parallel(self, tp_size=1, ep_size=1, cp_size=1):
+    def setup_model_parallel(self, tp_size=1, ep_size=1, cp_size=1, sequence_parallel=False):
         """Initialize model parallel with given configuration.
 
         Args:
@@ -615,6 +615,10 @@ def setup_model_parallel(self, tp_size=1, ep_size=1, cp_size=1):
             bf16=True,
             params_dtype=torch.bfloat16,
             add_bias_linear=False,
+            tensor_model_parallel_size=tp_size,
+            expert_model_parallel_size=ep_size,
+            context_parallel_size=cp_size,
+            sequence_parallel=sequence_parallel and tp_size > 1,
         )
 
     def new_router(self, **kwargs):
@@ -627,14 +631,15 @@ def new_router(self, **kwargs):
 
     @pytest.mark.internal
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+    @pytest.mark.parametrize("sequence_parallel", [True, False])
     @pytest.mark.parametrize("aux_loss_type", ["aux_loss", "seq_aux_loss", "global_aux_loss"])
     @pytest.mark.parametrize(
         "tp_size,ep_size,cp_size", [(8, 1, 1), (4, 2, 1), (1, 1, 8), (2, 1, 4), (2, 2, 2)]
     )
-    def test_padding_mask_removes_padding_tokens(self, aux_loss_type, tp_size, ep_size, cp_size):
+    def test_padding_mask_removes_padding_tokens(self, aux_loss_type, tp_size, ep_size, cp_size, sequence_parallel):
         """Test that padding tokens are correctly excluded from aux loss calculation."""
         # Initialize model parallel with given configuration
-        self.setup_model_parallel(tp_size=tp_size, ep_size=ep_size, cp_size=cp_size)
+        self.setup_model_parallel(tp_size=tp_size, ep_size=ep_size, cp_size=cp_size, sequence_parallel=sequence_parallel)
 
         try:
             clear_aux_losses_tracker()

From 13676a9f7e033e230017c48f1e35c6f3348c3b5e Mon Sep 17 00:00:00 2001
From: HaochenYuan <haocheny@nvidia.com>
Date: Tue, 23 Dec 2025 01:49:49 -0800
Subject: [PATCH 28/32] fix linting

---
 tests/unit_tests/transformer/moe/test_aux_loss.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/tests/unit_tests/transformer/moe/test_aux_loss.py b/tests/unit_tests/transformer/moe/test_aux_loss.py
index 319fbbea5c3..f5726777383 100644
--- a/tests/unit_tests/transformer/moe/test_aux_loss.py
+++ b/tests/unit_tests/transformer/moe/test_aux_loss.py
@@ -636,10 +636,14 @@ def new_router(self, **kwargs):
     @pytest.mark.parametrize(
         "tp_size,ep_size,cp_size", [(8, 1, 1), (4, 2, 1), (1, 1, 8), (2, 1, 4), (2, 2, 2)]
     )
-    def test_padding_mask_removes_padding_tokens(self, aux_loss_type, tp_size, ep_size, cp_size, sequence_parallel):
+    def test_padding_mask_removes_padding_tokens(
+        self, aux_loss_type, tp_size, ep_size, cp_size, sequence_parallel
+    ):
         """Test that padding tokens are correctly excluded from aux loss calculation."""
         # Initialize model parallel with given configuration
-        self.setup_model_parallel(tp_size=tp_size, ep_size=ep_size, cp_size=cp_size, sequence_parallel=sequence_parallel)
+        self.setup_model_parallel(
+            tp_size=tp_size, ep_size=ep_size, cp_size=cp_size, sequence_parallel=sequence_parallel
+        )
 
         try:
             clear_aux_losses_tracker()

From 6cf4ce6dace60b188554dbdb1d730cce323dda97 Mon Sep 17 00:00:00 2001
From: HaochenYuan <haocheny@nvidia.com>
Date: Tue, 23 Dec 2025 02:35:09 -0800
Subject: [PATCH 29/32] fix error

---
 tests/unit_tests/a2a_overlap/test_schedule_chunk_1f1b.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/unit_tests/a2a_overlap/test_schedule_chunk_1f1b.py b/tests/unit_tests/a2a_overlap/test_schedule_chunk_1f1b.py
index ad27187b63e..32f37c43d06 100644
--- a/tests/unit_tests/a2a_overlap/test_schedule_chunk_1f1b.py
+++ b/tests/unit_tests/a2a_overlap/test_schedule_chunk_1f1b.py
@@ -195,7 +195,7 @@ def test_1f1b_schedule_model_chunk_with_padding_mask(
         # Re-initialize model parallel with the specified configuration
         Utils.destroy_model_parallel()
         Utils.initialize_model_parallel(
-            tensor_model_parallel_size=parallel_config["tp"],
+            tensor_model_parallel_size=tp_size,
             pipeline_model_parallel_size=1,
             expert_model_parallel_size=4,
             expert_tensor_parallel_size=1,
@@ -212,8 +212,8 @@ def test_1f1b_schedule_model_chunk_with_padding_mask(
         # create TransformerConfig
         extra_kwargs = {
             "moe_token_dispatcher_type": dispatcher_type,
-            "tensor_model_parallel_size": parallel_config["tp"],
-            "sequence_parallel": parallel_config["tp"] > 1,
+            "tensor_model_parallel_size": tp_size,
+            "sequence_parallel": tp_size > 1,
         }
         if dispatcher_type == "flex":
             extra_kwargs["moe_flex_dispatcher_backend"] = "deepep"

From 0c8159b1636b3108a4da8f4378765d2bc43b4dc5 Mon Sep 17 00:00:00 2001
From: HaochenYuan <haocheny@nvidia.com>
Date: Tue, 23 Dec 2025 02:39:41 -0800
Subject: [PATCH 30/32] add copyright

---
 tests/test_utils/python_scripts/recipe_parser.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/test_utils/python_scripts/recipe_parser.py b/tests/test_utils/python_scripts/recipe_parser.py
index e26d04d6f20..d09dfa7a622 100644
--- a/tests/test_utils/python_scripts/recipe_parser.py
+++ b/tests/test_utils/python_scripts/recipe_parser.py
@@ -1,3 +1,4 @@
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 import copy
 import itertools
 import logging

From a0e159106047f193c416ec32268cd481b8a8193c Mon Sep 17 00:00:00 2001
From: HaochenYuan <haocheny@nvidia.com>
Date: Tue, 23 Dec 2025 03:09:06 -0800
Subject: [PATCH 31/32] fix error

---
 tests/unit_tests/a2a_overlap/test_schedule_chunk_1f1b.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/unit_tests/a2a_overlap/test_schedule_chunk_1f1b.py b/tests/unit_tests/a2a_overlap/test_schedule_chunk_1f1b.py
index 32f37c43d06..f5d2ba1779d 100644
--- a/tests/unit_tests/a2a_overlap/test_schedule_chunk_1f1b.py
+++ b/tests/unit_tests/a2a_overlap/test_schedule_chunk_1f1b.py
@@ -186,7 +186,7 @@ def test_1f1b_schedule_model_chunk(self, mtp_layers, dispatcher_type, fp8_flag,
     @pytest.mark.parametrize("layers", [[2, 1], [1, 1]])
     @pytest.mark.parametrize("tp_size", [1, 2, 4, 8])
     def test_1f1b_schedule_model_chunk_with_padding_mask(
-        self, dispatcher_type, layers, parallel_config
+        self, dispatcher_type, layers, tp_size
     ):
         """
         Verifies all-to-all overlap optimization with padding_mask produces

From 84344df2458eccbbee1401d6152d3bfc9567147e Mon Sep 17 00:00:00 2001
From: HaochenYuan <haocheny@nvidia.com>
Date: Tue, 23 Dec 2025 07:17:32 -0800
Subject: [PATCH 32/32] fix linting

---
 tests/unit_tests/a2a_overlap/test_schedule_chunk_1f1b.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/tests/unit_tests/a2a_overlap/test_schedule_chunk_1f1b.py b/tests/unit_tests/a2a_overlap/test_schedule_chunk_1f1b.py
index f5d2ba1779d..6c59dd3f9e3 100644
--- a/tests/unit_tests/a2a_overlap/test_schedule_chunk_1f1b.py
+++ b/tests/unit_tests/a2a_overlap/test_schedule_chunk_1f1b.py
@@ -185,9 +185,7 @@ def test_1f1b_schedule_model_chunk(self, mtp_layers, dispatcher_type, fp8_flag,
     @pytest.mark.parametrize("dispatcher_type", get_valid_token_dispatcher_types())
     @pytest.mark.parametrize("layers", [[2, 1], [1, 1]])
     @pytest.mark.parametrize("tp_size", [1, 2, 4, 8])
-    def test_1f1b_schedule_model_chunk_with_padding_mask(
-        self, dispatcher_type, layers, tp_size
-    ):
+    def test_1f1b_schedule_model_chunk_with_padding_mask(self, dispatcher_type, layers, tp_size):
         """
         Verifies all-to-all overlap optimization with padding_mask produces
         the same results as the reference implementation with various TP/EP/CP combinations.