From aa3fabaa19948445f43150a457b18aacab242190 Mon Sep 17 00:00:00 2001 From: gberasmus87 Date: Fri, 11 Sep 2026 11:53:31 +1200 Subject: [PATCH] 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. --- python/freetoken/models/gemma4/attention.py | 9 +++++++-- python/freetoken/models/muse_glimmer/attention.py | 9 +++++++-- python/freetoken/models/qwen3_5_moe/attention.py | 9 +++++++-- 3 files changed, 21 insertions(+), 6 deletions(-) diff --git a/python/freetoken/models/gemma4/attention.py b/python/freetoken/models/gemma4/attention.py index 4f85255f0..abfbd1b5d 100644 --- a/python/freetoken/models/gemma4/attention.py +++ b/python/freetoken/models/gemma4/attention.py @@ -5,7 +5,7 @@ import torch from freetoken.attention import AttentionSpec from freetoken.core import get_global_ctx -from freetoken.layers import BaseOP, GemmaRMSNorm, LinearQKVMerged, LinearReplicated +from freetoken.layers import BaseOP, GemmaRMSNorm, LinearOProj, LinearQKVMerged from freetoken.layers.rotary import get_rope from freetoken.models.config import FullAttentionGroupConfig, SWAAttentionGroupConfig from freetoken.utils import nvtx_annotate @@ -40,7 +40,12 @@ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""): quant_config=config.quant, prefix=f"{prefix}.qkv_proj", ) - self.o_proj = LinearReplicated( + # Row-parallel, as every other family with a column-parallel qkv builds o_proj: + # each rank's attention output is its local head slice, so o_proj takes the sharded + # input dim and all-reduces the partial sums. At TP=1 this degenerates to the previous + # replicated behaviour exactly (div_even(x, 1) == x, all-reduce skipped), so it is a + # no-op today and correct whenever this family gains tensor parallelism. + self.o_proj = LinearOProj( self.q_dim, config.hidden_size, has_bias=False, quant_config=config.quant, prefix=f"{prefix}.o_proj", ) diff --git a/python/freetoken/models/muse_glimmer/attention.py b/python/freetoken/models/muse_glimmer/attention.py index 7c814d4c1..aedef62b5 100644 --- a/python/freetoken/models/muse_glimmer/attention.py +++ b/python/freetoken/models/muse_glimmer/attention.py @@ -5,7 +5,7 @@ import torch from freetoken.attention import AttentionSpec from freetoken.core import get_global_ctx -from freetoken.layers import BaseOP, GemmaRMSNorm, LinearColParallelMerged, LinearReplicated +from freetoken.layers import BaseOP, GemmaRMSNorm, LinearColParallelMerged, LinearOProj from freetoken.layers.rotary import get_rope from freetoken.models.config import SWAAttentionGroupConfig from freetoken.utils import nvtx_annotate @@ -46,7 +46,12 @@ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""): config.hidden_size, self._qkvg_split, has_bias=False, quant_config=config.quant, prefix=f"{prefix}.qkvg_proj", ) - self.o_proj = LinearReplicated( + # Row-parallel, as every other family with a column-parallel qkv builds o_proj: + # each rank's attention output is its local head slice, so o_proj takes the sharded + # input dim and all-reduces the partial sums. At TP=1 this degenerates to the previous + # replicated behaviour exactly (div_even(x, 1) == x, all-reduce skipped), so it is a + # no-op today and correct whenever this family gains tensor parallelism. + self.o_proj = LinearOProj( self.qo_attn_dim, config.hidden_size, has_bias=False, quant_config=config.quant, prefix=f"{prefix}.o_proj", ) diff --git a/python/freetoken/models/qwen3_5_moe/attention.py b/python/freetoken/models/qwen3_5_moe/attention.py index 12a91bb3e..fde694489 100644 --- a/python/freetoken/models/qwen3_5_moe/attention.py +++ b/python/freetoken/models/qwen3_5_moe/attention.py @@ -4,7 +4,7 @@ import torch from freetoken.core import get_global_ctx -from freetoken.layers import BaseOP, GemmaRMSNorm, LinearColParallelMerged, LinearReplicated +from freetoken.layers import BaseOP, GemmaRMSNorm, LinearColParallelMerged, LinearOProj from freetoken.layers.rotary import get_rope from freetoken.utils import nvtx_annotate @@ -56,7 +56,12 @@ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""): else None ), ) - self.o_proj = LinearReplicated( + # Row-parallel, as every other family with a column-parallel qkv builds o_proj: + # each rank's attention output is its local head slice, so o_proj takes the sharded + # input dim and all-reduces the partial sums. At TP=1 this degenerates to the previous + # replicated behaviour exactly (div_even(x, 1) == x, all-reduce skipped), so it is a + # no-op today and correct whenever this family gains tensor parallelism. + self.o_proj = LinearOProj( self.qo_attn_dim, config.hidden_size, has_bias=False, quant_config=config.quant, prefix=f"{prefix}.o_proj", )