Skip to content

fix(models): build o_proj row-parallel in the remaining column-parallel families - #440

Open
gberasmus87 wants to merge 1 commit into
FlashML-org:mainfrom
gberasmus87:oproj-row-parallel-remaining
Open

fix(models): build o_proj row-parallel in the remaining column-parallel families#440
gberasmus87 wants to merge 1 commit into
FlashML-org:mainfrom
gberasmus87:oproj-row-parallel-remaining

Conversation

@gberasmus87

Copy link
Copy Markdown

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
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 #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.

…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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant