add UT

Author: Haochen Yuan

megatron/core/transformer/moe/moe_utils.py

@@ -34,6 +34,70 @@
 _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,

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,
@@ -276,16 +277,12 @@ def _apply_aux_loss(
         if aux_loss_coeff == 0:
             return probs
 
-        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_tokens = padding_mask.sum()
-        else:
-            tokens_per_expert = routing_map.sum(dim=0)
-            num_tokens = routing_map.shape[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
@@ -304,7 +301,7 @@ def _apply_aux_loss(
         )
         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.item() if isinstance(num_tokens, torch.Tensor) else num_tokens
+            valid_token_count=num_tokens
         )
         return probs
 
@@ -330,21 +327,17 @@ def _apply_seq_aux_loss(
 
         scores_for_aux_loss = scores_for_aux_loss.reshape(seq_length, -1)
 
-        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, self.config.num_moe_experts).reshape(seq_length, -1)
-            # Apply mask to routing_map for token counting
-            routing_map_masked = routing_map.reshape(seq_length, -1) & mask_expanded
-            tokens_per_expert = routing_map_masked.sum(dim=0)
-            # Count valid tokens per sequence
-            num_valid_tokens_per_seq = padding_mask_reshaped.sum(dim=0)  # [bsz]
-            total_num_tokens = num_valid_tokens_per_seq.sum() * self.tp_cp_group.size()
-        else:
-            tokens_per_expert = routing_map.reshape(seq_length, -1).sum(dim=0)
-            total_num_tokens = seq_length * self.tp_cp_group.size()
-            padding_mask_for_loss = None
+        # 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
@@ -365,8 +358,7 @@ def _apply_seq_aux_loss(
         )
         # Calculate valid token count: seq_length for each batch element
         if padding_mask is not None:
-            num_valid_tokens = padding_mask.sum()
-            valid_token_count = num_valid_tokens.item() if isinstance(num_valid_tokens, torch.Tensor) else num_valid_tokens
+            valid_token_count = padding_mask.sum()
         else:
             valid_token_count = seq_length * bsz
 
@@ -385,16 +377,12 @@ def _apply_global_aux_loss(
         if global_aux_loss_coeff == 0:
             return probs
 
-        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_tokens = padding_mask.sum()
-        else:
-            tokens_per_expert = routing_map.sum(dim=0)
-            num_tokens = scores_for_aux_loss.shape[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
@@ -422,7 +410,7 @@ def _apply_global_aux_loss(
             global_aux_loss,
             "global_load_balancing_loss",
             self.tp_dp_cp_group,
-            valid_token_count=num_tokens.item() if isinstance(num_tokens, torch.Tensor) else num_tokens,
+            valid_token_count=num_tokens,
         )
         return probs
 
@@ -433,7 +421,7 @@ def attach_and_log_load_balancing_loss(
         aux_loss: torch.Tensor,
         aux_loss_name: str,
         reduce_group: torch.distributed.ProcessGroup,
-        valid_token_count: Optional[int] = None,
+        valid_token_count: Optional[Union[int, torch.Tensor]] = None,
     ):
         """Attach aux loss function to activation and add to logging.
         
@@ -443,7 +431,8 @@ def attach_and_log_load_balancing_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, optional): Number of valid tokens excluding padding tokens.
+            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

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()