[diffusion] fix: keep a LoRA unmerged under merge_mode auto when merging would round its update away

docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx

@@ -756,9 +756,13 @@ sglang serve \
 ```
 
 `auto` merges an adapter into ordinary resident weights to avoid per-step LoRA
-matmuls, but keeps the dynamic path for FSDP-sharded weights where a full
-gather can increase peak memory. Use `dynamic` when one resident server must
-switch repeatedly between base and LoRA output.
+matmuls. It keeps the dynamic path in two cases. One is FSDP-sharded weights,
+where a full gather can increase peak memory. The other is an adapter whose
+update would mostly round away in a BF16 merge. Distilled adapters such as the
+Larry Turbo LoRA sit below a BF16 ulp of most weights, and a merge drops
+77–97% of their update in typical layers. Use `dynamic` when one resident
+server must switch repeatedly between base and LoRA output, and `merge` to
+force a merge.
 
 Use the filename, scale, and request schedule from the table together. The
 4-evaluation LightX2V recipe is the more aggressive latency/quality tradeoff.

python/sglang/multimodal_gen/runtime/layers/lora/linear.py

@@ -448,6 +448,37 @@ def compute_merged_weight(self) -> torch.Tensor:
         self._merge_lora_into_data(work, self.lora_weights_list)
         return work.to("cpu", dtype=target_dtype)
 
+    @torch.no_grad()
+    def merge_rounding_norms(self, max_rows: int = 1024) -> tuple[float, float] | None:
+        """Squared norms of the part of the active LoRA update that merging
+        into this layer's weight dtype would round away, and of the update,
+        over an even sample of rows. Nothing is written; None when there is
+        nothing to round or the adapter has no plain 2-D form here."""
+        lora_list = self._active_lora_list()
+        if self.disable_lora or not lora_list:
+            return None
+        src, _ = self._materialized_weight_src()
+        if not src.is_floating_point() or src.dtype == torch.float32:
+            return None
+        weight = src.reshape(-1, src.shape[-1])
+        device = get_local_torch_device()
+        rows = torch.linspace(0, weight.shape[0] - 1, min(max_rows, weight.shape[0]))
+        rows = rows.round().long().unique()
+        base = weight[rows.to(weight.device)].to(device, torch.float32)
+        delta = torch.zeros_like(base)
+        for lora_A, lora_B, _, strength, rank, alpha, _ in lora_list:
+            a = self.slice_lora_a_weights(lora_A.to(device, torch.float32))
+            b = self.slice_lora_b_weights(lora_B.to(device, torch.float32))
+            if not isinstance(b, torch.Tensor) or a.dim() > 2 or b.dim() > 2:
+                return None
+            scale = strength
+            if alpha is not None and rank is not None and alpha != rank:
+                scale *= alpha / rank
+            delta += scale * (b[rows.to(device)] @ a)
+        exact = base + delta
+        lost = exact.to(src.dtype).to(torch.float32) - exact
+        return float(lost.square().sum()), float(delta.square().sum())
+
     def install_merged_weight(
         self, merged: torch.Tensor, base_view: torch.Tensor
     ) -> None:

python/sglang/multimodal_gen/runtime/pipelines_core/lora/pipeline.py

@@ -57,6 +57,26 @@
 
 logger = init_logger(__name__)
 
+# merge_mode "auto" keeps an adapter unmerged when merging it into the base
+# weights would round away more than this share of its update (by norm), i.e.
+# when the merge would destroy most of it. Distilled LoRAs sit below a BF16 ulp
+# of the weights (MiniMax-H3 Turbo loses 75%); style LoRAs keep most of theirs
+# (Wan Arcane-Jinx loses 14%, a Z-Image anime LoRA 31%) and would pay the
+# dynamic path's cost, up to +50% per step on small models, for little
+AUTO_MERGE_MAX_ROUNDING_LOSS = 0.5
+
+
+def merge_rounding_loss(lora_layers: dict) -> float | None:
+    """Share of the attached LoRA update (by norm) that merging it into the
+    layers' weight dtype would round away; None when nothing can be measured."""
+    lost = total = 0.0
+    for layer in lora_layers.values():
+        norms = layer.merge_rounding_norms()
+        if norms is not None:
+            lost += norms[0]
+            total += norms[1]
+    return (lost / total) ** 0.5 if total > 0 else None
+
 
 def _swap_peft_swiglu_fc1_lora_b(
     source_name: str, target_name: str, weight: torch.Tensor
@@ -187,6 +207,8 @@ def __init__(self, *args, **kwargs) -> None:
         self.lora_layers_critic = {}
         self.lora_layers_transformer_2 = {}
         self.is_lora_merged = {}
+        # merge_mode "auto": rounding loss per (module, paths, strengths, alphas)
+        self.auto_merge_rounding_loss = {}
         self.lora_initialized = False
         self.lora_rank = None
         self.lora_alpha = None
@@ -1189,6 +1211,20 @@ def set_lora(
                 # Apply LoRA to modules for this target
                 for module_name, lora_layers_dict in target_modules:
                     effective_merge_weights = merge_weights_by_module[module_name]
+                    if (
+                        effective_merge_weights
+                        and merge_mode == "auto"
+                        and len(tgt_nicknames) == 1
+                        and self._auto_merge_rounds_away(
+                            module_name,
+                            lora_layers_dict,
+                            tgt_nicknames,
+                            tgt_paths,
+                            rank,
+                            tgt_strengths,
+                        )
+                    ):
+                        effective_merge_weights = False
                     count = None
                     if not effective_merge_weights and not adapter_updated:
                         count = self._reactivate_cached_dynamic_lora_layers(
@@ -1253,6 +1289,48 @@ def set_lora(
             merge_mode,
         )
 
+    def _auto_merge_rounds_away(
+        self,
+        module_name: str,
+        lora_layers: dict[str, BaseLayerWithLoRA],
+        nicknames: list[str],
+        paths: list[str | None],
+        rank: int,
+        strengths: list[float],
+    ) -> bool:
+        """Whether merge_mode "auto" should keep this adapter unmerged on the
+        module. It is attached unmerged to measure, which is where it stays."""
+        key = (
+            module_name,
+            tuple(paths),
+            tuple(strengths),
+            tuple(self.loaded_adapter_alphas.get(n) for n in nicknames),
+        )
+        if key not in self.auto_merge_rounding_loss:
+            self._apply_lora_to_layers(
+                lora_layers,
+                nicknames,
+                paths,
+                rank,
+                strengths,
+                clear_existing=True,
+                merge_weights=False,
+            )
+            self.auto_merge_rounding_loss[key] = merge_rounding_loss(lora_layers)
+        loss = self.auto_merge_rounding_loss[key]
+        if loss is None or loss <= AUTO_MERGE_MAX_ROUNDING_LOSS:
+            return False
+        if rank == 0:
+            logger.info(
+                "LoRA %s stays unmerged on %s: merging it into the base weights "
+                "would round away %.0f%% of its update. Pass "
+                "--lora-merge-mode merge to merge anyway.",
+                ",".join(str(p) for p in paths),
+                module_name,
+                100 * loss,
+            )
+        return True
+
     def _merge_via_cache(self, name, layer, merge_cache) -> None:
         """Merge one layer through the cache instead of in place.
 

python/sglang/multimodal_gen/runtime/server_args/server_args.py

@@ -3093,9 +3093,11 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
             choices=LORA_MERGE_MODES,
             default=ServerArgs.lora_merge_mode,
             help=(
-                "How LoRA is applied: auto keeps static merge for regular weights "
-                "and uses dynamic LoRA for FSDP-sharded weights to avoid full-gather; "
-                "merge always merges into base weights; dynamic always applies LoRA at forward time."
+                "How LoRA is applied: auto merges into regular weights, but uses "
+                "dynamic LoRA for FSDP-sharded weights (to avoid a full gather) and "
+                "for adapters a merge would mostly round away (more than 50%% of "
+                "the update, e.g. distilled LoRAs); merge always merges into base "
+                "weights; dynamic always applies LoRA at forward time."
             ),
         )
         parser.add_argument(

python/sglang/multimodal_gen/test/unit/test_lora_auto_merge_rounding.py

@@ -0,0 +1,98 @@
+"""merge_mode "auto" keeps a LoRA unmerged when merging would round its update away.
+
+Distilled adapters (MiniMax-H3 Turbo) carry updates below a BF16 ulp of the base
+weights, so a merge into BF16 drops most of them. These check the estimate against
+an actual merge on CPU, and the decision the pipeline takes from it.
+"""
+
+from types import SimpleNamespace
+
+import torch
+from torch import nn
+
+from sglang.multimodal_gen.runtime.layers.lora.linear import LinearWithLoRA
+from sglang.multimodal_gen.runtime.pipelines_core.lora.pipeline import (
+    AUTO_MERGE_MAX_ROUNDING_LOSS,
+    LoRAPipeline,
+    merge_rounding_loss,
+)
+
+
+def _layer(base: torch.Tensor, a: torch.Tensor, b: torch.Tensor, strength: float):
+    linear = nn.Linear(base.shape[1], base.shape[0], bias=False, dtype=base.dtype)
+    with torch.no_grad():
+        linear.weight.copy_(base)
+    layer = LinearWithLoRA(linear, lora_rank=a.shape[0], lora_alpha=a.shape[0])
+    layer.set_lora_weights(
+        a, b, strength=strength, clear_existing=True, merge_weights=False
+    )
+    return layer
+
+
+def _inputs(update_scale: float, seed: int = 0):
+    g = torch.Generator().manual_seed(seed)
+    base = torch.randn(96, 64, generator=g).to(torch.bfloat16)
+    a = torch.randn(8, 64, generator=g)
+    b = update_scale * torch.randn(96, 8, generator=g)
+    return base, a, b
+
+
+def test_estimate_matches_the_rounding_of_an_actual_merge() -> None:
+    base, a, b = _inputs(1e-3)
+    layer = _layer(base, a, b, strength=0.5)
+    lost_sq, delta_sq = layer.merge_rounding_norms(max_rows=10**6)
+    exact = base.float() + 0.5 * (b @ a)
+    layer.merge_lora_weights()
+    merged = layer.base_layer.weight.detach().float()
+    torch.testing.assert_close(
+        torch.tensor(lost_sq), (merged - exact).square().sum(), rtol=1e-3, atol=0
+    )
+    torch.testing.assert_close(
+        torch.tensor(delta_sq), (0.5 * (b @ a)).square().sum(), rtol=1e-4, atol=0
+    )
+
+
+def test_update_below_a_bf16_ulp_is_mostly_rounded_away() -> None:
+    base, a, b = _inputs(1e-5)
+    loss = merge_rounding_loss({"l": _layer(base, a, b, strength=1.0)})
+    assert loss is not None and loss > 0.9
+
+
+def test_update_above_a_bf16_ulp_survives() -> None:
+    base, a, b = _inputs(0.1)
+    loss = merge_rounding_loss({"l": _layer(base, a, b, strength=1.0)})
+    assert loss is not None and loss < 0.05
+
+
+def test_fp32_base_has_nothing_to_round() -> None:
+    base, a, b = _inputs(1e-5)
+    layer = _layer(base.float(), a, b, strength=1.0)
+    assert layer.merge_rounding_norms() is None
+    assert merge_rounding_loss({"l": layer}) is None
+
+
+def test_auto_decision_measures_once_and_follows_the_threshold() -> None:
+    calls = []
+    norms = {"lossy": (0.81, 1.0), "fine": (1e-4, 1.0)}
+
+    def run(kind: str) -> bool:
+        layers = {"l": SimpleNamespace(merge_rounding_norms=lambda: norms[kind])}
+        stub = SimpleNamespace(
+            auto_merge_rounding_loss={},
+            loaded_adapter_alphas={},
+            _apply_lora_to_layers=lambda *args, **kwargs: calls.append(kwargs),
+        )
+        first = LoRAPipeline._auto_merge_rounds_away(
+            stub, "transformer", layers, ["a"], [kind], 0, [1.0]
+        )
+        second = LoRAPipeline._auto_merge_rounds_away(
+            stub, "transformer", layers, ["a"], [kind], 0, [1.0]
+        )
+        assert first == second
+        return first
+
+    assert run("lossy") is True  # 90% of the update's norm lost
+    assert run("fine") is False
+    # measured once per configuration, attached unmerged to measure
+    assert len(calls) == 2 and all(c["merge_weights"] is False for c in calls)
+    assert AUTO_MERGE_MAX_ROUNDING_LOSS == 0.5