[Feat] opt qwen image model load use ColumnParallelLinear replace ReplicatedLinear (#3875)

Signed-off-by: rongfu.leng <lenronfu@gmail.com>
Co-authored-by: Peter Pan <Peter.Pan@daocloud.io>

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"),