[cosmos3] Bitwise parity across LoRA updates: unmerged weight sync + engine-side adapter replay

miles/backends/fsdp_utils/diffusion_update_weight_utils.py

@@ -312,17 +312,39 @@ def update_bucket_weights(
                 ray.get(ref)
 
 
+_FORWARD_DTYPES = {"fp32": torch.float32, "bf16": torch.bfloat16, "fp16": torch.float16}
+
+
+def _strip_peft_wrappers(name: str) -> str:
+    """base_model.model.X.base_layer.weight -> X.weight (see the comment in
+    ``_update_component_weights`` for the wrapper layout)."""
+    name = name.replace(".base_layer", "")
+    if name.startswith("base_model.model."):
+        name = name[len("base_model.model.") :]
+    return name
+
+
 # TODO: update weights only for sgl-d LoRA params
 class DiffusionUpdateWeightFromTensorLoRA(DiffusionUpdateWeightFromTensor):
     """LoRA-aware updater: merges adapters into base before pushing to rollout.
 
     The rollout engine has no LoRA layers — it receives standard weight keys
     like ``transformer_blocks.0.attn.to_q.weight``.  We compute ``W_base + αBA/r``
     on the fly during sync (no in-place mutation of the FSDP model).
+
+    With ``--lora-unmerged-weight-sync`` the base weights ship untouched and the
+    adapters ship alongside as ``<layer>.lora_A.weight`` / ``<layer>.lora_B.weight``
+    / ``<layer>.lora_scaling`` (A/B pre-rounded to the trainer's forward dtype —
+    the same rounding FSDP's mixed-precision gather applies in the train
+    forward). An engine-side patch stores them and adds the adapter GEMMs in
+    peft's exact op order, keeping the two forwards bitwise comparable after
+    updates: ``GEMM(W + αBA/r)`` is not bitwise equal to
+    ``GEMM(W) + GEMM_B(GEMM_A(x))·s``, so merged sync caps parity at one step.
     """
 
     def __init__(self, args, models):
         super().__init__(args, models)
+        self._unmerged = getattr(args, "lora_unmerged_weight_sync", False)
         # Per-component LoRA index: component -> {param name -> (A, B, scaling)}.
         self._lora_index: dict[str, dict[str, tuple]] = {}
         for component, model in self.models.items():
@@ -336,7 +358,24 @@ def __init__(self, args, models):
                             module.scaling[adapter],
                         )
             self._lora_index[component] = index
-            logger.info(f"LoRA weight sync [{component}]: {len(index)} mergeable layers")
+            mode = "unmerged (adapter tensors)" if self._unmerged else "merged"
+            logger.info(f"LoRA weight sync [{component}]: {len(index)} layers, mode={mode}")
+
+    def _iter_adapter_tensors(self, lora_index: dict[str, tuple]):
+        """Yield (name, tensor) adapter entries for unmerged sync."""
+        forward_dtype = _FORWARD_DTYPES[self.args.diffusion_forward_dtype]
+        for base_name, (A, B, s) in lora_index.items():
+            stripped = _strip_peft_wrappers(base_name)
+            assert stripped.endswith(".weight")
+            prefix = stripped[: -len(".weight")]
+            # Round exactly as the train forward sees the adapters: FSDP gathers
+            # the fp32 masters at forward dtype (elementwise cast, shard-order
+            # invariant), so `.to(forward_dtype)` reproduces those bits.
+            yield f"{prefix}.lora_A.weight", self._gather_full(A.weight.detach()).to(forward_dtype).contiguous()
+            yield f"{prefix}.lora_B.weight", self._gather_full(B.weight.detach()).to(forward_dtype).contiguous()
+            # fp64 so the engine recovers the exact python float peft multiplies
+            # by (`... * self.scaling`); fp32 could round e.g. alpha/r = 10/3.
+            yield f"{prefix}.lora_scaling", torch.tensor([float(s)], dtype=torch.float64, device="cuda")
 
     def _gather_full(self, t: torch.Tensor) -> torch.Tensor:
         t = t.cuda()
@@ -366,7 +405,7 @@ def _update_component_weights(self, target_module: str, model: torch.nn.Module)
                     async_op=True,
                 ).to_local()
 
-            if name in lora_index:
+            if name in lora_index and not self._unmerged:
                 # Merge LoRA for this layer on the fly instead of pre-computing
                 # all 720 deltas up front: Qwen-Image's MLP + attn deltas total
                 # tens of GB at peak — here only one delta is resident at a time.
@@ -387,9 +426,7 @@ def _update_component_weights(self, target_module: str, model: torch.nn.Module)
             #
             # ``.base_layer`` is the inner wrapper (lora.Linear.base_layer);
             # ``base_model.model.`` is PeftModel.base_model (=LoraModel) .model.
-            sglang_d_param_name = name.replace(".base_layer", "")
-            if sglang_d_param_name.startswith("base_model.model."):
-                sglang_d_param_name = sglang_d_param_name[len("base_model.model.") :]
+            sglang_d_param_name = _strip_peft_wrappers(name)
 
             sz = param.numel() * param.element_size()
             if bucket and bucket_size + sz >= self.args.update_weight_buffer_size:
@@ -403,6 +440,15 @@ def _update_component_weights(self, target_module: str, model: torch.nn.Module)
                 t = param.wait() if hasattr(param, "wait") else param
                 verify_pairs.append((sglang_d_param_name, t.detach().cpu().contiguous()))
 
+        if self._unmerged:
+            for adapter_name, tensor in self._iter_adapter_tensors(lora_index):
+                sz = tensor.numel() * tensor.element_size()
+                if bucket and bucket_size + sz >= self.args.update_weight_buffer_size:
+                    self.wait_and_update_bucket_weights(bucket, target_module)
+                    bucket, bucket_size = [], 0
+                bucket.append((adapter_name, tensor))
+                bucket_size += sz
+
         if bucket:
             self.wait_and_update_bucket_weights(bucket, target_module)
 

miles/backends/sglang_diffusion_utils/monkey_patches/patch_cosmos3_bitwise.py

@@ -32,6 +32,14 @@
     train side is single-sample by construction (``compute_noise_pred``
     asserts ``not cfg_batching``). This also removes the uncond text padding
     (11 -> 29) that batch-2 forced.
+  * LoRA runs as adapter GEMMs instead of a weight merge. The trainer's peft
+    forward is ``base(x) + lora_B(lora_A(x))·s`` (three GEMMs); a merged
+    ``GEMM(W + sBA)`` rounds differently, so merged sync caps parity at the
+    first step (B starts at 0). With ``--lora-unmerged-weight-sync`` the
+    trainer ships base weights untouched plus per-layer A/B/scaling tensors;
+    this side intercepts them at the weight-sync loader and replays peft's
+    exact op sequence per target (to_qkv slices for add_q/k/v, to_out for
+    to_add_out).
 """
 
 from __future__ import annotations
@@ -47,6 +55,7 @@ def apply() -> None:
     _patch_silu_and_mul_eager()
     _patch_qk_norm_rope_split_eager()
     _patch_cfg_sequential()
+    _patch_lora_adapter_intercept()
 
 
 def _force_torch_sdpa_backend() -> None:
@@ -96,17 +105,124 @@ def _forward(self, x: torch.Tensor):
         if not logged:
             logged = True
             print(f"[cosmos3_bitwise] unfused MergedColumnParallelLinear active: slices={sizes}", flush=True)
+        lora = getattr(self, "_miles_lora", None)
         outs = []
         offset = 0
-        for size in sizes:
+        for idx, size in enumerate(sizes):
             bias = self.bias[offset : offset + size] if self.bias is not None else None
-            outs.append(F.linear(x, self.weight[offset : offset + size], bias))
+            out = F.linear(x, self.weight[offset : offset + size], bias)
+            if lora is not None and idx in lora:
+                out = out + _lora_term(x, *lora[idx])
+            outs.append(out)
             offset += size
         return torch.cat(outs, dim=-1), None
 
     MergedColumnParallelLinear.forward = _forward
 
 
+def _lora_term(x: torch.Tensor, A: torch.Tensor, B: torch.Tensor, s: float) -> torch.Tensor:
+    """peft vanilla LoRA (0.18.x), op for op: ``lora_B(lora_A(x)) * scaling``.
+
+    The trainer runs it under bf16 autocast on FSDP-gathered bf16 adapters; the
+    shipped A/B are pre-rounded to that dtype, so both sides execute the same
+    two bf16 GEMMs and the same elementwise multiply. The base output is added
+    by the caller as ``base + term`` — same order as peft's ``result + ...``.
+    """
+    return F.linear(F.linear(x, A), B) * s
+
+
+def _attach_lora_adapters(module, adapters: dict[str, dict[str, torch.Tensor]]) -> None:
+    """Store shipped adapter tensors on their target submodules.
+
+    ``adapters`` maps a diffusers-style layer prefix (e.g.
+    ``layers.5.self_attn.add_q_proj``) to its ``.lora_A.weight`` /
+    ``.lora_B.weight`` / ``.lora_scaling`` tensors. The module's own
+    param-name mapper resolves the prefix to the sgl-d parameter — including
+    the merge index for slices of a fused param (add_q/k/v -> to_qkv slots
+    0/1/2; to_add_out -> to_out, no index).
+    """
+    from sglang.multimodal_gen.runtime.post_training.weights_updater import (
+        _build_module_weight_name_mapper,
+    )
+
+    map_name = _build_module_weight_name_mapper(module)
+    for prefix, parts in adapters.items():
+        missing = {".lora_A.weight", ".lora_B.weight", ".lora_scaling"} - set(parts)
+        if missing:
+            raise RuntimeError(f"cosmos3_bitwise LoRA sync: incomplete adapter for {prefix!r}: missing {missing}")
+        mapped, slot = map_name(f"{prefix}.weight") if map_name is not None else (f"{prefix}.weight", None)
+        target = module.get_submodule(mapped[: -len(".weight")])
+        device = next(target.parameters()).device
+        # clone(): the shipped tensors are views into the CUDA-IPC flattened
+        # bucket, whose storage the sender reclaims after the update returns.
+        A = parts[".lora_A.weight"].to(device).clone()
+        B = parts[".lora_B.weight"].to(device).clone()
+        s = float(parts[".lora_scaling"].item())
+        registry = getattr(target, "_miles_lora", None)
+        if registry is None:
+            registry = {}
+            target._miles_lora = registry
+        registry[slot] = (A, B, s)
+        if slot is None:
+            _wrap_linear_instance_with_lora(target)
+
+
+def _wrap_linear_instance_with_lora(target) -> None:
+    """Instance-level wrap for non-fused targets (RowParallelLinear to_out):
+    add the adapter term after the complete base output (bias included), the
+    position peft adds it at."""
+    if getattr(target, "_miles_lora_wrapped", False):
+        return
+    orig_forward = target.forward
+
+    def forward(x):
+        out, out_bias = orig_forward(x)
+        A, B, s = target._miles_lora[None]
+        return out + _lora_term(x, A, B, s), out_bias
+
+    target.forward = forward
+    target._miles_lora_wrapped = True
+
+
+def _patch_lora_adapter_intercept() -> None:
+    """Consume `<layer>.lora_A/lora_B/lora_scaling` tensors from weight sync.
+
+    The trainer's --lora-unmerged-weight-sync ships base weights untouched
+    plus per-layer adapter tensors (see DiffusionUpdateWeightFromTensorLoRA).
+    sgl-d's loader would warn-and-drop these unknown names, so split them out
+    before it runs and attach them to the resolved target modules. No-op when
+    the trainer syncs merged weights.
+
+    A layer's three parts arrive in separate calls — the sender flushes one
+    flattened bucket per dtype (A/B at forward dtype, scaling fp64) and may
+    also split across buffer-size flushes — so partial adapters are buffered
+    until complete.
+    """
+    from sglang.multimodal_gen.runtime.post_training import weights_updater
+
+    orig_load = weights_updater._load_weights_into_module
+    pending: dict[str, dict[str, torch.Tensor]] = {}
+
+    def _load(module, weights_iter):
+        base_entries = []
+        for name, weight in weights_iter:
+            for suffix in (".lora_A.weight", ".lora_B.weight", ".lora_scaling"):
+                if name.endswith(suffix):
+                    pending.setdefault(name[: -len(suffix)], {})[suffix] = weight
+                    break
+            else:
+                base_entries.append((name, weight))
+        complete = {prefix: parts for prefix, parts in pending.items() if len(parts) == 3}
+        if complete:
+            _attach_lora_adapters(module, complete)
+            for prefix in complete:
+                del pending[prefix]
+            print(f"[cosmos3_bitwise] attached {len(complete)} LoRA adapters (unmerged weight sync)", flush=True)
+        return orig_load(module, iter(base_entries))
+
+    weights_updater._load_weights_into_module = _load
+
+
 def _patch_silu_and_mul_eager() -> None:
     from sglang.multimodal_gen.runtime.layers.activation import SiluAndMul
 

miles/utils/arguments.py

@@ -1028,6 +1028,19 @@ def add_lora_arguments(parser):
                     "(requires matching sglang-d LoRAPipeline support)."
                 ),
             )
+            parser.add_argument(
+                "--lora-unmerged-weight-sync",
+                action="store_true",
+                default=False,
+                help=(
+                    "Ship base weights unmerged plus lora_A/lora_B (+scaling) tensors so the "
+                    "rollout engine applies LoRA as adapter GEMMs in peft's exact op order "
+                    "instead of merging W+BA. Keeps train/rollout forwards bitwise comparable "
+                    "after weight updates (a merged GEMM is not bitwise equal to base+adapter). "
+                    "Requires an engine-side patch that consumes the adapter tensors (see "
+                    "monkey_patches.patch_cosmos3_bitwise)."
+                ),
+            )
             return parser
 
         def add_ema_arguments(parser):