From c349f62839ebf3fafc13c7a06a9be6670363103c Mon Sep 17 00:00:00 2001 From: "rongfu.leng" Date: Wed, 27 May 2026 17:02:17 +0800 Subject: [PATCH] [Feat] opt qwen image model load use ColumnParallelLinear replace ReplicatedLinear (#3875) Signed-off-by: rongfu.leng Co-authored-by: Peter Pan --- .../models/qwen_image/qwen_image_transformer.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) 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"),