diff --git a/miles/backends/fsdp_utils/diffusion_update_weight_utils.py b/miles/backends/fsdp_utils/diffusion_update_weight_utils.py index baba2f06..8e83285c 100644 --- a/miles/backends/fsdp_utils/diffusion_update_weight_utils.py +++ b/miles/backends/fsdp_utils/diffusion_update_weight_utils.py @@ -312,6 +312,18 @@ 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. @@ -319,10 +331,20 @@ class DiffusionUpdateWeightFromTensorLoRA(DiffusionUpdateWeightFromTensor): 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 ``.lora_A.weight`` / ``.lora_B.weight`` + / ``.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) diff --git a/miles/backends/sglang_diffusion_utils/monkey_patches/patch_cosmos3_bitwise.py b/miles/backends/sglang_diffusion_utils/monkey_patches/patch_cosmos3_bitwise.py index 9314e90c..f91cae70 100644 --- a/miles/backends/sglang_diffusion_utils/monkey_patches/patch_cosmos3_bitwise.py +++ b/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 `.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 diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index dc57ddd2..87b9e6cd 100644 --- a/miles/utils/arguments.py +++ b/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):