fix(models): build o_proj row-parallel in the remaining column-parallel families - #440
Open
gberasmus87 wants to merge 1 commit into
Open
fix(models): build o_proj row-parallel in the remaining column-parallel families#440gberasmus87 wants to merge 1 commit into
gberasmus87 wants to merge 1 commit into
Conversation
…el families Seven families pair a column-parallel qkv with `LinearOProj` (llama, qwen2, qwen3, qwen3_moe, mistral, gpt_oss, minimax_m2). Three pair one with `LinearReplicated` instead: gemma4, qwen3_5_moe and muse_glimmer. That pairing is wrong under TP>1 in three ways at once. The qkv projection shards by head -- `LinearQKVMerged` computes `div_even(num_qo_heads, tp_info.size)`, `LinearColParallelMerged` shards its output sizes -- so a rank's attention output is its local head slice rather than the full width. `o_proj` therefore has to take the sharded input dim, and the partial sums need an all-reduce. `LinearReplicated` keeps the full weight, expects the unsharded input, and reduces nothing. It also fails quietly: a missing all-reduce leaves each rank holding a partial sum that still decodes to fluent-looking text, so a token-level smoke test does not catch it. None of the three is reachable today -- all three weight loaders raise `NotImplementedError(... supports TP=1 only)` -- so this is a latent trap rather than a live bug, and the change is a no-op for every current deployment. `LinearOProj` degenerates to exactly the previous replicated behaviour at TP=1: `div_even(x, 1) == x`, and the all-reduce is skipped when `tp_size == 1`. It adds no new constraint from calling `get_tp_info()` in `__init__` either, since each of these constructors already reaches it one or two lines up through its qkv class. The families left alone are correct as they stand: glm5_next, glm_moe_dsa and minimax_m3 replicate their q/kv projections too, so replicated `o_proj` is consistent there, and glm4_moe uses `LinearDF11`. Same change FlashML-org#429 makes for qwen4_exp, which @gdevenyi independently validated at TP=2 (GSM8K 97.00%, unchanged). Doing the remaining three in one pass so the next family to gain tensor parallelism does not rediscover this. Verified: `LinearOProj` and `LinearReplicated` forwards are bit-identical at TP=1 at each family's real attention dimensions, and each family's test suite shows the same pass/fail set with and without the change.
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.
fix(models): build o_proj row-parallel in the remaining column-parallel families
Seven families pair a column-parallel qkv with
LinearOProj(llama, qwen2,qwen3, qwen3_moe, mistral, gpt_oss, minimax_m2). Three pair one with
LinearReplicatedinstead: gemma4, qwen3_5_moe and muse_glimmer.That pairing is wrong under TP>1 in three ways at once. The qkv projection
shards by head --
LinearQKVMergedcomputesdiv_even(num_qo_heads, tp_info.size),LinearColParallelMergedshards its output sizes -- so a rank'sattention output is its local head slice rather than the full width.
o_projtherefore has to take the sharded input dim, and the partial sums need an
all-reduce.
LinearReplicatedkeeps the full weight, expects the unshardedinput, and reduces nothing.
It also fails quietly: a missing all-reduce leaves each rank holding a partial
sum that still decodes to fluent-looking text, so a token-level smoke test does
not catch it.
None of the three is reachable today -- all three weight loaders raise
NotImplementedError(... supports TP=1 only)-- so this is a latent trap ratherthan a live bug, and the change is a no-op for every current deployment.
LinearOProjdegenerates to exactly the previous replicated behaviour at TP=1:div_even(x, 1) == x, and the all-reduce is skipped whentp_size == 1. It addsno new constraint from calling
get_tp_info()in__init__either, since eachof these constructors already reaches it one or two lines up through its qkv
class.
The families left alone are correct as they stand: glm5_next, glm_moe_dsa and
minimax_m3 replicate their q/kv projections too, so replicated
o_projisconsistent there, and glm4_moe uses
LinearDF11.Same change #429 makes for qwen4_exp, which @gdevenyi independently validated at
TP=2 (GSM8K 97.00%, unchanged). Doing the remaining three in one pass so the
next family to gain tensor parallelism does not rediscover this.
Verified:
LinearOProjandLinearReplicatedforwards are bit-identical at TP=1at each family's real attention dimensions, and each family's test suite shows
the same pass/fail set with and without the change.