From e7224e96816b8a1a2e0e07a6a63d5fa6f30c2d1d Mon Sep 17 00:00:00 2001
From: Fan Lin <49047353+linfann@users.noreply.github.com>
Date: Wed, 21 Jan 2026 18:17:16 +0800
Subject: [PATCH] [diffusion] fix: fix the LoRA weights mismatch caused by
 weights packing (#17355)

---
 .../multimodal_gen/configs/models/dits/zimage.py       | 10 ++++++++++
 .../multimodal_gen/runtime/layers/lora/linear.py       |  6 ++++++
 .../runtime/pipelines_core/lora_pipeline.py            |  8 ++++----
 3 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/python/sglang/multimodal_gen/configs/models/dits/zimage.py b/python/sglang/multimodal_gen/configs/models/dits/zimage.py
index de3275852..bef5e948e 100644
--- a/python/sglang/multimodal_gen/configs/models/dits/zimage.py
+++ b/python/sglang/multimodal_gen/configs/models/dits/zimage.py
@@ -38,6 +38,16 @@ class ZImageArchConfig(DiTArchConfig):
         default_factory=lambda: {
             r"(.*)\.feed_forward\.w1\.weight$": (r"\1.feed_forward.w13.weight", 0, 2),
             r"(.*)\.feed_forward\.w3\.weight$": (r"\1.feed_forward.w13.weight", 1, 2),
+            r"(.*)\.feed_forward\.w1\.(lora_A|lora_B)$": (
+                r"\1.feed_forward.w13.\2",
+                0,
+                2,
+            ),
+            r"(.*)\.feed_forward\.w3\.(lora_A|lora_B)$": (
+                r"\1.feed_forward.w13.\2",
+                1,
+                2,
+            ),
         }
     )
 
diff --git a/python/sglang/multimodal_gen/runtime/layers/lora/linear.py b/python/sglang/multimodal_gen/runtime/layers/lora/linear.py
index 2ef6f4d70..ebeabef94 100644
--- a/python/sglang/multimodal_gen/runtime/layers/lora/linear.py
+++ b/python/sglang/multimodal_gen/runtime/layers/lora/linear.py
@@ -90,6 +90,8 @@ class BaseLayerWithLoRA(nn.Module):
                     self.lora_alpha / self.lora_rank  # type: ignore
                 )  # type: ignore
             delta = delta * self.strength
+            if delta.dim() > 2:
+                delta = delta.reshape(-1, delta.shape[-1])
             out, output_bias = self.base_layer(x)
             return out + delta, output_bias
         else:
@@ -171,6 +173,8 @@ class BaseLayerWithLoRA(nn.Module):
             if self.lora_alpha is not None and self.lora_rank is not None:
                 if self.lora_alpha != self.lora_rank:
                     lora_delta = lora_delta * (self.lora_alpha / self.lora_rank)
+            if lora_delta.dim() > 2:
+                lora_delta = lora_delta.reshape(-1, lora_delta.shape[-1])
             data += lora_strength * lora_delta
 
     @torch.no_grad()
@@ -468,6 +472,8 @@ class LinearWithLoRA(BaseLayerWithLoRA):
                     self.lora_alpha / self.lora_rank  # type: ignore
                 )  # type: ignore
             delta = delta * self.strength
+            if delta.dim() > 2:
+                delta = delta.reshape(-1, delta.shape[-1])
             # nn.Linear.forward() returns a single tensor, not a tuple
             out = self.base_layer(x)
             return out + delta
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py
index 6239153a9..f4e0565a7 100644
--- a/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py
+++ b/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py
@@ -560,17 +560,17 @@ class LoRAPipeline(ComposedPipelineBase):
             name, _, _ = lora_param_names_mapping_fn(name)
             # HF-format (LoRA) -> SGLang-dit-format
             target_name, merge_index, num_params_to_merge = param_names_mapping_fn(name)
-            # for (in_dim, r) @ (r, out_dim), we only merge (r, out_dim * n) where n is the number of linear layers to fuse
+            # for fuse B(out_dim, r) @ A(r, in_dim) -> (N, out_dim, r) @ (N, r, in_dim)
             # see param mapping in HunyuanVideoArchConfig
-            if merge_index is not None and "lora_B" in name:
+            if merge_index is not None:
                 to_merge_params[target_name][merge_index] = weight
                 if len(to_merge_params[target_name]) == num_params_to_merge:
-                    # cat at output dim according to the merge_index order
                     sorted_tensors = [
                         to_merge_params[target_name][i]
                         for i in range(num_params_to_merge)
                     ]
-                    weight = torch.cat(sorted_tensors, dim=1)
+                    # Use stack instead of cat because it needs to be compatible with TP.
+                    weight = torch.stack(sorted_tensors, dim=0)
                     del to_merge_params[target_name]
                 else:
                     continue