fix(moe): broadcast padding mask for expert bias counts

megatron/core/transformer/moe/router.py

@@ -738,7 +738,7 @@ def _apply_expert_bias(
         if self.enable_expert_bias and torch.is_grad_enabled():
             with torch.no_grad():
                 if padding_mask is not None:
-                    routing_map = routing_map & (~padding_mask)
+                    routing_map = routing_map & (~padding_mask).unsqueeze(-1)
                 self.local_tokens_per_expert += routing_map.sum(dim=0)
 
     def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = None):

tests/unit_tests/transformer/moe/test_routers.py

@@ -1,6 +1,7 @@
 # Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
 
 
+import dataclasses
 from typing import cast
 
 import pytest
@@ -179,6 +180,62 @@ def test_router_with_padding_mask(self):
             # Verify that probs for valid tokens are similar
             assert torch.equal(probs_valid_part, probs_without_mask)
 
+    @pytest.mark.internal
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+    @pytest.mark.parametrize("with_padding_mask", [False, True])
+    def test_expert_bias_token_counts_with_padding_mask(self, with_padding_mask):
+        """Test expert-bias counts in grad-enabled masked and unmasked forwards."""
+        config = dataclasses.replace(
+            self.transformer_config,
+            moe_router_enable_expert_bias=True,
+            moe_router_load_balancing_type="none",
+            moe_router_score_function="sigmoid",
+        )
+        submodules = get_submodules(
+            get_gpt_layer_local_submodules(
+                num_experts=config.num_moe_experts, moe_grouped_gemm=False
+            ).mlp
+        )
+        assert isinstance(submodules, MoESubmodules)
+        router = cast(Router, MoELayer(config, submodules).router).cuda().train()
+
+        seq_len = 5
+        batch_size = 2
+        hidden_states = torch.randn(
+            (seq_len, batch_size, config.hidden_size),
+            dtype=torch.bfloat16,
+            device="cuda",
+            requires_grad=True,
+        )
+        assert seq_len * batch_size != config.num_moe_experts
+        padding_mask = None
+        if with_padding_mask:
+            padding_mask = torch.zeros((seq_len, batch_size), dtype=torch.bool, device="cuda")
+            padding_mask[-2:, 0] = True
+
+        probs, routing_map = router(hidden_states, padding_mask=padding_mask)
+        valid_routing_map = routing_map
+        if padding_mask is not None:
+            valid_routing_map = routing_map[~padding_mask.reshape(-1)]
+
+        expected_tokens_per_expert = valid_routing_map.sum(dim=0)
+        torch.testing.assert_close(
+            router.local_tokens_per_expert,
+            expected_tokens_per_expert.to(router.local_tokens_per_expert.dtype),
+        )
+        expected_valid_tokens = seq_len * batch_size
+        if padding_mask is not None:
+            expected_valid_tokens -= padding_mask.sum().item()
+        assert (
+            router.local_tokens_per_expert.sum() == expected_valid_tokens * config.moe_router_topk
+        )
+
+        probs.sum().backward()
+        assert hidden_states.grad is not None
+        assert hidden_states.grad.isfinite().all()
+        assert router.weight.grad is not None
+        assert router.weight.grad.isfinite().all()
+
     @pytest.mark.internal
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
     def test_router_dtype(self):