From 2973a722725c58eabf691e998fe2fc5ede6da4b6 Mon Sep 17 00:00:00 2001 From: gberasmus87 Date: Thu, 10 Sep 2026 16:39:00 +1200 Subject: [PATCH] qwen4_exp: build o_proj row-parallel, as every other family does qwen4_exp is the only family that builds its attention o_proj as LinearReplicated; llama, gpt_oss and minimax_m2 all use LinearOProj. That is correct at TP=1 and wrong under TP>1 three ways at once: qkv_proj is column-parallel, so a rank's attention output is its local head slice rather than the full qo_attn_dim, o_proj therefore has to take the sharded input dim, and the partial sums need an all-reduce. LinearReplicated keeps the full [hidden, qo_attn_dim] 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. LinearOProj degenerates to exactly LinearReplicated at TP=1 -- div_even(x, 1) == x, and the all-reduce is skipped when tp_size == 1 -- so this is a no-op for main as it stands and only changes what #385 finds when the two meet. It adds no constraint from calling get_tp_info() in __init__ either, since the same constructor already reaches it two lines up through LinearColParallelMerged. The comment above the branch now describes what is built. Raised by @gdevenyi against the earlier form of this work in #392. --- python/freetoken/models/qwen4_exp/attention.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/python/freetoken/models/qwen4_exp/attention.py b/python/freetoken/models/qwen4_exp/attention.py index 66211c7c2..3bb6f1f04 100644 --- a/python/freetoken/models/qwen4_exp/attention.py +++ b/python/freetoken/models/qwen4_exp/attention.py @@ -19,7 +19,7 @@ import torch from freetoken.core import get_global_ctx -from freetoken.layers import BaseOP, GemmaPlusOneRMSNorm, LinearColParallelMerged, LinearReplicated +from freetoken.layers import BaseOP, GemmaPlusOneRMSNorm, LinearColParallelMerged, LinearOProj, LinearReplicated from freetoken.layers.rotary import get_rope from freetoken.utils import nvtx_annotate @@ -128,7 +128,11 @@ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = "") -> N config.hidden_size, self._qkv_split, has_bias=False, quant_config=config.quant, prefix=f"{prefix}.qkv_proj", ) - self.o_proj = LinearReplicated( + # Row-parallel, as every other family builds o_proj: qkv_proj is column-parallel, so a + # rank's attention output is its local head slice and the partial sums need an + # all-reduce. At TP=1 this is LinearReplicated exactly -- div_even(x, 1) == x and the + # all-reduce is skipped -- so it is a no-op today and correct when #385 lands. + self.o_proj = LinearOProj( self.qo_attn_dim, config.hidden_size, has_bias=False, quant_config=config.quant, prefix=f"{prefix}.o_proj", )