DeepSpeed

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

Clamp capacity to num_tokens in MoE gating drop branches (complete #5353) (#8155) ## 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. --------- Signed-off-by: Ehsan Barkhordar <realbarkhordar@gmail.com> Co-authored-by: Masahiro Tanaka <mtanaka@anyscale.com> Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com>