diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py index c3f99f7f999..a2dada0c026 100644 --- a/megatron/core/transformer/moe/router.py +++ b/megatron/core/transformer/moe/router.py @@ -789,6 +789,20 @@ def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = N router_replay=self.router_replay, ) + # Dropless HybridEP consumes the sparse routing map directly, so exclude padding + # rows before dispatch. Other dispatchers retain their existing fixed-route + # assumptions until they support sparse routing maps end to end. + use_dropless_hybridep = ( + self.config.moe_token_dispatcher_type == "flex" + and self.config.moe_flex_dispatcher_backend == "hybridep" + and self.config.moe_expert_capacity_factor is None + and self.config.moe_expert_rank_capacity_factor is None + ) + if padding_mask is not None and use_dropless_hybridep: + valid_tokens = (~padding_mask).unsqueeze(-1) + probs = probs * valid_tokens + routing_map = routing_map & valid_tokens + # Apply token dropping to probs and routing_map. if self.config.moe_expert_capacity_factor is not None: probs, routing_map = apply_router_token_dropping( diff --git a/tests/unit_tests/transformer/moe/test_routers.py b/tests/unit_tests/transformer/moe/test_routers.py index 6a2ab7908aa..8fc63194c2c 100644 --- a/tests/unit_tests/transformer/moe/test_routers.py +++ b/tests/unit_tests/transformer/moe/test_routers.py @@ -135,9 +135,15 @@ def test_aux_loss(self): @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.""" + @pytest.mark.parametrize("router_fusion", [False, True]) + def test_router_with_padding_mask(self, router_fusion): + """Test that HybridEP excludes padding tokens from routing.""" + if router_fusion and not HAVE_ROUTER_FUSION: + pytest.skip("TE fused router ops not available") self.router = self.router.cuda() + self.router.config.moe_router_fusion = router_fusion + self.router.config.moe_token_dispatcher_type = "flex" + self.router.config.moe_flex_dispatcher_backend = "hybridep" seq_len = 32 batch_size = 2 hidden_size = self.router.config.hidden_size @@ -177,9 +183,50 @@ def test_router_with_padding_mask(self): self.router.config.num_moe_experts, ) + padding_rows = padding_mask.reshape(-1) + assert torch.count_nonzero(probs_with_mask[padding_rows]) == 0 + assert not routing_map_with_mask[padding_rows].any() + # 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( + "dispatcher,backend,capacity_factor,rank_capacity_factor", + [ + ("allgather", "deepep", None, None), + ("alltoall", "deepep", None, None), + ("flex", "deepep", None, None), + ("flex", "deepepv2", None, None), + ("flex", "hybridep", 1.0, None), + ("flex", "hybridep", None, 1.0), + ], + ) + def test_padding_mask_preserves_routes_outside_dropless_hybridep( + self, dispatcher, backend, capacity_factor, rank_capacity_factor + ): + """Only dropless HybridEP may consume a sparse route map.""" + self.router = self.router.cuda() + self.router.config.moe_token_dispatcher_type = dispatcher + self.router.config.moe_flex_dispatcher_backend = backend + self.router.config.moe_expert_capacity_factor = capacity_factor + self.router.config.moe_expert_rank_capacity_factor = rank_capacity_factor + hidden_states = torch.randn( + (16, 2, self.router.config.hidden_size), device="cuda", dtype=torch.bfloat16 + ) + padding_mask = torch.zeros((16, 2), dtype=torch.bool, device="cuda") + padding_mask[8:, :] = True + + with torch.no_grad(): + probs_with_mask, routing_map_with_mask = self.router( + hidden_states, padding_mask=padding_mask + ) + probs_without_mask, routing_map_without_mask = self.router(hidden_states) + + torch.testing.assert_close(probs_with_mask, probs_without_mask) + assert torch.equal(routing_map_with_mask, routing_map_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])