diff --git a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py index dac724162..5e4d526d8 100644 --- a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py +++ b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py @@ -716,15 +716,17 @@ class QwenImageTransformerBlock(nn.Module): self.attention_head_dim = attention_head_dim # Image processing modules. - # Modulation linear is kept full precision (quant_config=None) — it + # Modulation linear is kept unquantized (quant_config=None) — it # produces shift/scale/gate values that are precision-sensitive - # (see #2728). + # (see #2728). Use column TP with gather_output=True so weights are + # sharded while downstream modulation still receives full [B, 6 * dim]. self.img_mod = nn.Sequential( nn.SiLU(), - ReplicatedLinear( + ColumnParallelLinear( dim, 6 * dim, bias=True, + gather_output=True, return_bias=False, quant_config=None, prefix=_join_prefix(prefix, "img_mod.1"), @@ -751,10 +753,11 @@ class QwenImageTransformerBlock(nn.Module): # Text processing modules. self.txt_mod = nn.Sequential( nn.SiLU(), - ReplicatedLinear( + ColumnParallelLinear( dim, 6 * dim, bias=True, + gather_output=True, return_bias=False, quant_config=None, prefix=_join_prefix(prefix, "txt_mod.1"),