Repository navigation
Conversation
…not split the KV heads Since #48859 (v5.19), `tp_size=ep_size=N` also applies the TP plan, so a MoE with fewer KV heads than N (gemma-4-26B-A4B-it, Qwen3.5-35B-A3B at 8 ranks) crashes at the first attention view; up to v5.18 expert parallelism left attention whole. Under EP, when the KV heads cannot be split by `tp_size`, drop the attention rules from the TP plan and warn, and keep the rest of the plan. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
Contributor
CI recapDashboard: View test results in Grafana |
Contributor
Author
|
After discussing with @ArthurZucker and @IlyasMoutawwakil, the correct fix at this stage is #49633 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #49628.
Fix. In
resolve_parallel_plans, under expert parallelism, whenget_kv_heads_per_rankreports thattp_sizecannot split the KV heads, the attention rules are dropped from the TP plan (and frommodel.tp_plan, whichgenerateuses to size the static cache) with a warning. Attention stays replicated as in v5.18; embeddings,lm_head, the dense MLPs and the experts are sharded as onmain. Models whose heads divide are unchanged, and so is pure TP (ep_size=1).Tests.
test_expert_parallel_keeps_unsplittable_attention_replicatedresolves the plans of a tiny Qwen3-MoE with 2 and 4 KV heads attp_size=ep_size=4;test_expert_parallel_forward_with_fewer_kv_heads_than_ranksruns the 2-KV-head model on 4 gloo ranks and compares the logits with a single process.cc @ArthurZucker @vasqu @3outeille @IlyasMoutawwakil