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

[ebarkhordar]

Problem

torch.topk(x, k=capacity, dim=0) over the token dimension requires capacity <= num_tokens. In the MoE gating code capacity is ceil(num_tokens / num_experts * capacity_factor * k), so whenever capacity_factor * k > num_experts the capacity exceeds the number of tokens and the topk call raises RuntimeError: selected index k out of range.

PR #5353 ("Ensure capacity does not exceed number of tokens") added the guard capacity = min(capacity, num_tokens), but only to top1gating's no-drop branch. The two drop branches that actually feed capacity into torch.topk(..., dim=0) were left unguarded:

• topkgating, drop_policy='probs': torch.topk(topk_masked_gates, k=capacity, dim=0) (sharded_moe.py:410)
• top1gating, drop_tokens=True: _top_idx(mask1_rand, capacity) → torch.topk(source, k=capacity, dim=0) (sharded_moe.py:176/248)

drop_policy='position' and top2gating use torch.lt rather than torch.topk, so they are unaffected and are left untouched.

Fix

Apply the same capacity <= num_tokens clamp #5353 introduced to the two drop branches. Reducing capacity to num_tokens is a no-op whenever capacity already fits, and it cannot drop any token that should route (a per-expert dispatch buffer never needs more slots than there are tokens).

Invariant

capacity passed to torch.topk(..., k=capacity, dim=0) over the token dimension is <= num_tokens on every gating path. Enumerated the torch.topk(..., dim=0) call sites in sharded_moe.py with ast (not grep): exactly two, both listed above; torch.topk(gates, k=k, dim=1) at :392 is over the expert dimension with k <= num_experts and is not involved.

Verification

Reproduced and fixed in a clean python:3.10-slim Docker container at HEAD, pip install -e ., torch==2.5.1+cpu (confirmed sharded_moe.__file__ resolves to the installed source):

• On the unpatched branches, topkgating(logits[8,2], k=2, capacity_factor=2, drop_policy='probs') raises RuntimeError: selected index k out of range, and top1gating(logits[8,2], capacity_factor=4, drop_tokens=True) raises the same error via _top_idx.
• With the clamp, both calls return; the dispatch buffer's capacity dimension equals num_tokens, and no routed token is dropped.
• Two regression tests added to tests/unit/moe/test_moe.py (one per branch); both fail on the unpatched source and pass with the fix. The pre-existing test_moe.py gating tests still pass.

Not verified: I did not measure whether a real training configuration in the wild sets capacity_factor * k > num_experts; the argument here is that the crash path exists and that #5353 already established capacity > num_tokens as a condition worth guarding.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 737050a869

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Preserve dispatch capacity when clamping top-k

When tensor-model parallelism is enabled, MOELayer.forward later calls drop_tokens on the capacity dimension, and drop_tokens asserts that this dimension is divisible by the TP world size (deepspeed/moe/mappings.py:64-65). For top-1 with drop_tokens=True, capacity_factor > num_experts, and a local token count that is not TP-divisible, this line sets the dispatch capacity to num_tokens; _top_idx succeeds, but the subsequent TP drop still fails. Clamp only the k passed to _top_idx (or repad the dispatch capacity) instead of reusing the clamped value for _one_hot_to_float.

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Keep the clamp out of the position policy

When callers use topkgating(..., drop_policy='position') with a configured capacity or min_capacity greater than the number of tokens, this clamp changes the returned dispatch capacity even though that branch never passes capacity to torch.topk(..., dim=0). Previously, for example, min_capacity=8 with 4 tokens returned a capacity width of 8; now it returns 4, violating the min-capacity/padding behavior and removing padding that downstream code may require. Move the clamp into the probs branch or use a separate top-k selection limit.

Useful? React with 👍 / 👎.

[goransh-buh]

Nice, straightforward implementation.