Clamp capacity to num_tokens in MoE gating drop branches (complete #5353)

deepspeed/moe/sharded_moe.py

@@ -245,7 +245,10 @@ def top1gating(logits: Tensor,
     assert logits.shape[
         0] >= min_capacity, "No. of tokens (batch-size) should be greater than min_capacity. Either set min_capacity to 0 or increase your batch size."
 
-    top_idx = _top_idx(mask1_rand, capacity)
+    # torch.topk(..., dim=0) cannot select more rows than the token dimension holds, so bound the
+    # selection without shrinking the dispatch capacity the buffers are sized from.
+    selection_capacity = min(capacity, torch.tensor(mask1.size(0)).to(capacity.device))
+    top_idx = _top_idx(mask1_rand, selection_capacity)
 
     new_mask1 = mask1 * torch.zeros_like(mask1).scatter_(0, top_idx, 1)
     mask1 = new_mask1
@@ -415,7 +418,10 @@ def topkgating(
 
         if drop_policy == 'probs':
             topk_masked_gates = torch.zeros_like(gates).scatter(1, top_idx, top_gate)
-            _, capacity_indices = torch.topk(topk_masked_gates, k=capacity, dim=0, sorted=False)
+            # k cannot exceed the token dimension being selected over, but capacity itself still
+            # sizes the dispatch buffer below.
+            selection_capacity = min(capacity, torch.tensor(gates.size(0)).to(capacity.device))
+            _, capacity_indices = torch.topk(topk_masked_gates, k=selection_capacity, dim=0, sorted=False)
             capacity_mask = torch.zeros_like(gates, dtype=torch.bool).scatter_(0, capacity_indices, True)
             mask &= capacity_mask
             locations = torch.cumsum(mask, dim=0) - 1

tests/unit/moe/test_moe.py

@@ -383,6 +383,42 @@ def check_equal(logits, cap, sparse_truth, res):
         check_equal(logits3, 4, sec_sparse, dispatch_res)
 
 
+# The drop branches select at most num_tokens rows with torch.topk(..., dim=0), while the
+# configured capacity still sizes the dispatch width.
+def test_topkgating_probs_capacity_exceeds_num_tokens():
+    # s=8, e=2, k=2, capacity_factor=2 -> capacity = ceil(8/2 * (2*2)) = 16 > 8.
+    num_tokens, num_experts, k = 8, 2, 2
+    logits = torch.randn(num_tokens, num_experts)
+    _, _, dispatch_mask, _ = topkgating(logits, k, capacity_factor=2.0, min_capacity=0, drop_policy='probs')
+    assert dispatch_mask.shape[-1] == 16
+    assert int(dispatch_mask.sum()) == num_tokens * k
+
+
+def test_top1gating_drop_capacity_exceeds_num_tokens():
+    # s=8, e=2, capacity_factor=4 -> capacity = ceil(8/2 * 4) = 16 > 8. This is the
+    # top1gating drop branch, which reaches torch.topk via _top_idx.
+    num_tokens, num_experts = 8, 2
+    logits = torch.randn(num_tokens, num_experts)
+    _, _, dispatch_mask, _ = top1gating(logits, capacity_factor=4.0, min_capacity=0, drop_tokens=True, use_rts=False)
+    assert dispatch_mask.shape[-1] == 16
+    assert int(dispatch_mask.sum()) == num_tokens
+
+
+def test_topkgating_position_preserves_min_capacity():
+    # drop_policy='position' never selects over the token dimension, so min_capacity padding stays.
+    num_tokens, num_experts = 4, 2
+    logits = torch.randn(num_tokens, num_experts)
+    _, _, dispatch_mask, _ = topkgating(logits, 1, capacity_factor=1.0, min_capacity=8, drop_policy='position')
+    assert dispatch_mask.shape[-1] == 8
+
+
+def test_top1gating_preserves_tensor_parallel_capacity():
+    # s=3, e=4, capacity_factor=8 -> capacity = 6, which drop_tokens() needs divisible by tp=2.
+    logits = torch.randn(3, 4)
+    _, _, dispatch_mask, _ = top1gating(logits, capacity_factor=8.0, min_capacity=0, drop_tokens=True, use_rts=False)
+    assert dispatch_mask.shape[-1] == 6
+
+
 class TestExpertWeightGradWithZero(DistributedTest):
     world_size = 2