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

deepspeed/moe/sharded_moe.py

@@ -224,6 +224,10 @@ def top1gating(logits: Tensor,
             new_capacity = torch.ceil(new_capacity / tp).mul(tp).to(new_capacity.dtype)
         # Make sure the capacity value does not exceed the number of tokens.
         capacity = min(new_capacity, torch.tensor(mask1.size(0)).to(new_capacity.device))
+    else:
+        # Same guard for the drop branch: capacity feeds torch.topk(..., dim=0) over the token
+        # dimension in _top_idx below, which requires capacity <= num_tokens.
+        capacity = min(capacity, torch.tensor(mask1.size(0)).to(capacity.device))
 
     # Compute l_aux
     me = torch.mean(gates, dim=0)
@@ -403,6 +407,11 @@ def topkgating(
     if drop_tokens:
         # Calculate configured capacity and remove locations outside capacity from mask
         capacity = _capacity(gates, torch.tensor(capacity_factor * k), torch.tensor(min_capacity))
+        # Make sure the capacity value does not exceed the number of tokens. The drop_policy=='probs'
+        # branch below selects tokens with torch.topk(..., k=capacity, dim=0) over the token
+        # dimension, which requires capacity <= num_tokens; #5353 added this same clamp to
+        # top1gating's no-drop branch but not to the drop branches.
+        capacity = min(capacity, torch.tensor(gates.size(0)).to(capacity.device))
         # update mask and locations by capacity
 
         if drop_policy == 'probs':

tests/unit/moe/test_moe.py

@@ -354,6 +354,34 @@ def check_equal(logits, cap, sparse_truth, res):
         check_equal(logits3, 4, sec_sparse, dispatch_res)
 
 
+# #5353 clamped capacity to the number of tokens only in top1gating's no-drop branch. The drop
+# branches feed capacity to torch.topk(..., dim=0) over the token dimension, which requires
+# capacity <= num_tokens, so when capacity_factor * k > num_experts the capacity computed as
+# ceil(num_tokens / num_experts * capacity_factor * k) exceeds num_tokens and torch.topk raises
+# "selected index k out of range". These single-process gating checks fail on the unpatched
+# branches and pass once the same clamp is applied there.
+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')
+    # Capacity is clamped to num_tokens, and no routed token is spuriously dropped
+    # (every token keeps all k experts because capacity now covers all tokens).
+    assert dispatch_mask.shape[-1] == num_tokens
+    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] == num_tokens
+    # top1 routes each token to exactly one expert; with capacity == num_tokens none are dropped.
+    assert int(dispatch_mask.sum()) == num_tokens
+
+
 class TestExpertWeightGradWithZero(DistributedTest):
     world_size = 2