From 737050a869eca9c66998ff68c34b4b977846f5fe Mon Sep 17 00:00:00 2001 From: Ehsan Barkhordar Date: Sat, 18 Jul 2026 07:52:17 +0000 Subject: [PATCH 1/2] Clamp capacity to num_tokens in MoE gating drop branches torch.topk(x, k=capacity, dim=0) over the token dimension requires capacity <= num_tokens. PR #5353 added this clamp to top1gating's no-drop branch but not to the drop branches, where capacity feeds torch.topk: topkgating drop_policy='probs' and top1gating's drop path (via _top_idx) raise 'selected index k out of range' when capacity_factor * k > num_experts. Apply the same clamp to both drop branches and add a regression test per branch. Signed-off-by: Ehsan Barkhordar --- deepspeed/moe/sharded_moe.py | 9 +++++++++ tests/unit/moe/test_moe.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 37 insertions(+) diff --git a/deepspeed/moe/sharded_moe.py b/deepspeed/moe/sharded_moe.py index d2a6c089e8e7..a1f1c0272baa 100644 --- a/deepspeed/moe/sharded_moe.py +++ b/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': diff --git a/tests/unit/moe/test_moe.py b/tests/unit/moe/test_moe.py index 6283007d3e8e..34e8a889cd5c 100644 --- a/tests/unit/moe/test_moe.py +++ b/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 From 0c691a52a6e00461dfd59eff07a840aaecb30d89 Mon Sep 17 00:00:00 2001 From: Ehsan Barkhordar Date: Thu, 30 Jul 2026 17:55:57 +0000 Subject: [PATCH 2/2] Bound the top-k selection instead of the dispatch capacity The previous commit clamped capacity itself, which also shrank the dispatch width: drop_tokens() asserts that width is divisible by the TP size, and drop_policy='position' relies on it for min_capacity padding. Clamp only the k passed to torch.topk(..., dim=0) and leave capacity as configured. Adds a regression test for each of the two behaviours. Co-authored-by: Masahiro Tanaka Signed-off-by: Ehsan Barkhordar --- deepspeed/moe/sharded_moe.py | 19 ++++++++----------- tests/unit/moe/test_moe.py | 30 +++++++++++++++++++----------- 2 files changed, 27 insertions(+), 22 deletions(-) diff --git a/deepspeed/moe/sharded_moe.py b/deepspeed/moe/sharded_moe.py index a1f1c0272baa..629ed468df1c 100644 --- a/deepspeed/moe/sharded_moe.py +++ b/deepspeed/moe/sharded_moe.py @@ -224,10 +224,6 @@ 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) @@ -249,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 @@ -407,16 +406,14 @@ 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': 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 diff --git a/tests/unit/moe/test_moe.py b/tests/unit/moe/test_moe.py index 34e8a889cd5c..4b8ad67326bb 100644 --- a/tests/unit/moe/test_moe.py +++ b/tests/unit/moe/test_moe.py @@ -354,20 +354,14 @@ 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. +# 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') - # 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 dispatch_mask.shape[-1] == 16 assert int(dispatch_mask.sum()) == num_tokens * k @@ -377,11 +371,25 @@ def test_top1gating_drop_capacity_exceeds_num_tokens(): 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 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