From f9d53710af023248a29d1c2db676e6c64a784bd8 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Sat, 8 Aug 2026 10:41:37 +0000 Subject: [PATCH 1/3] =?UTF-8?q?feat(cosmos3):=20bitwise=20train/rollout=20?= =?UTF-8?q?parity=20=E2=80=94=20fp32=20time=5Fembedder=20island=20+=20cosm?= =?UTF-8?q?os3=5Fbitwise=20patch=20group?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Train side: gather time_embedder at fp32 via the family FSDPParallelPlan's param_dtype_patterns (diffusers keeps it in _keep_in_fp32_modules and sgl-d loads it fp32; a blanket bf16 gather silently downgraded it), dedupe the per-token timestep rows to sgl-d's M=1 GEMM shape under autocast-off, and reroute diffusers RMSNorm through F.rms_norm (fp32-through-mul). Rollout side (`--rollout-patch-group cosmos3_bitwise`): pin TORCH_SDPA, F.rms_norm RMSNorm, unfuse MergedColumnParallelLinear into per-slice GEMMs, eager SiluAndMul, split qk-norm+rope path, and sequential batch-1 CFG — the exact op sequence the diffusers train forward runs, at equal precision. Verified: 103/103 aligned tensor pairs bit-exact across both CFG branches, both towers, first/last denoise steps (embed_tokens -> proj_out). Co-authored-by: Cursor --- miles/backends/fsdp_utils/configs/cosmos3.py | 62 +++++- .../models/diffusers/cosmos3/parallel_plan.py | 11 +- .../monkey_patches/__init__.py | 7 + .../monkey_patches/patch_cosmos3_bitwise.py | 192 ++++++++++++++++++ ...ffusion-grpo-cosmos3-pickscore-t2i-5gpu.sh | 1 + 5 files changed, 270 insertions(+), 3 deletions(-) create mode 100644 miles/backends/sglang_diffusion_utils/monkey_patches/patch_cosmos3_bitwise.py diff --git a/miles/backends/fsdp_utils/configs/cosmos3.py b/miles/backends/fsdp_utils/configs/cosmos3.py index 6a4f33eb..faf5e32e 100644 --- a/miles/backends/fsdp_utils/configs/cosmos3.py +++ b/miles/backends/fsdp_utils/configs/cosmos3.py @@ -51,6 +51,9 @@ class Cosmos3TrainPipelineConfig(TrainPipelineConfig): # asserts it); never batch the CFG branches. cfg_batching = False lora_target_modules = ["add_q_proj", "add_k_proj", "add_v_proj", "to_add_out"] + # time_embedder gathers at fp32 via the family FSDPParallelPlan + # (models/diffusers/cosmos3/parallel_plan.py); rollout parity patches ship + # as the `cosmos3_bitwise` group, selected with --rollout-patch-group. @classmethod def validate_args(cls, args) -> None: @@ -181,10 +184,65 @@ def postprocess_model_after_materialize(self, model: torch.nn.Module) -> None: param.requires_grad_(False) # sglang-d casts the fp32 timestep sinusoid to the MLP weight dtype - # before linear_1; diffusers feeds it through as-is, which crashes on - # the fp32/bf16 mismatch under FSDP mixed precision. + # before linear_1 (`t_freq.to(w_dtype)`); mirror that exactly. With the + # fp32 pattern in the family FSDPParallelPlan the weights gather at + # fp32, so this keeps the sinusoid at fp32 like sglang-d's time_embedder. def _cast_to_weight_dtype(module, args): dtype = module.linear_1.weight.dtype return tuple(a.to(dtype) if torch.is_tensor(a) else a for a in args) model.time_embedder.register_forward_pre_hook(_cast_to_weight_dtype) + _wrap_time_embedder_row_dedup(model.time_embedder) + _patch_diffusers_rmsnorm_fp32_through_mul() + + +def _wrap_time_embedder_row_dedup(time_embedder: torch.nn.Module) -> None: + """Collapse identical sinusoid rows before the timestep MLP, expand after. + + sglang-d runs the timestep MLP once per request (GEMM M=1) and broadcasts + the embedding over tokens; diffusers expands the timestep per token first + (M=390 for a 480x480 clip). cuBLAS fp32 GEMMs are not bitwise M-invariant + (measured: linear_2 4096->4096 differs between M=1 and M=2), so per-token + rows can never bit-match the rollout engine. Deduplicating is numerically + exact — the rows are byte-identical copies — and reproduces sglang-d's + compute shape. Gradients are unchanged up to the usual expand/sum autograd. + """ + orig_forward = time_embedder.forward + + def forward(x, *args, **kwargs): + # Autocast off: the trainer's bf16 autocast would cast the fp32-gathered + # weights back to bf16 at the matmul boundary, undoing the parallel + # plan's fp32 island. sgl-d runs this MLP at plain fp32 with no autocast. + with torch.autocast("cuda", enabled=False): + if torch.is_tensor(x) and x.ndim == 2 and x.shape[0] > 1 and torch.equal(x, x[:1].expand_as(x)): + out = orig_forward(x[:1], *args, **kwargs) + return out.expand(x.shape[0], *out.shape[1:]) + return orig_forward(x, *args, **kwargs) + + time_embedder.forward = forward + + +def _patch_diffusers_rmsnorm_fp32_through_mul() -> None: + """Raise diffusers RMSNorm to fp32-through-the-weight-mul via F.rms_norm. + + diffusers' eager RMSNorm rounds the normalized activations to bf16 BEFORE + multiplying the weight (two bf16 roundings); sglang-d keeps fp32 through + the weight mul and rounds once. Following the "never downgrade" rule the + train side comes up: route through torch's fused F.rms_norm (fp32 + accumulation, single rounding). The rollout patch group routes sglang-d's + RMSNorm through the same op, so both sides run identical kernels. + """ + from diffusers.models import normalization + + if getattr(normalization.RMSNorm, "_miles_fp32_through_mul", False): + return + + orig_forward = normalization.RMSNorm.forward + + def forward(self, hidden_states): + if self.weight is not None and self.bias is None: + return torch.nn.functional.rms_norm(hidden_states, self.dim, self.weight, self.eps) + return orig_forward(self, hidden_states) + + normalization.RMSNorm.forward = forward + normalization.RMSNorm._miles_fp32_through_mul = True diff --git a/miles/backends/fsdp_utils/models/diffusers/cosmos3/parallel_plan.py b/miles/backends/fsdp_utils/models/diffusers/cosmos3/parallel_plan.py index daec57f0..7c229b15 100644 --- a/miles/backends/fsdp_utils/models/diffusers/cosmos3/parallel_plan.py +++ b/miles/backends/fsdp_utils/models/diffusers/cosmos3/parallel_plan.py @@ -1,4 +1,13 @@ from miles.backends.fsdp_utils.models.parallel_plan import FSDPParallelPlan -FSDP_PARALLEL_PLAN = FSDPParallelPlan() +FSDP_PARALLEL_PLAN = FSDPParallelPlan( + param_dtype_patterns={ + # diffusers declares `_keep_in_fp32_modules = ["time_embedder"]` and + # sglang-d pins the same module to fp32 at load; a blanket FSDP bf16 + # gather silently downgraded it on the train side. Gather it at fp32 to + # restore the upstream precision contract (log-prob parity with the + # rollout engine's fp32 MLP). + "*time_embedder*": "fp32", + }, +) diff --git a/miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py b/miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py index e377f75e..0e39567a 100644 --- a/miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py +++ b/miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py @@ -55,6 +55,13 @@ def apply_sgld_monkey_patches() -> None: patch_qk_norm_rope.apply() +@register_rollout_patch_group("cosmos3_bitwise") +def apply_cosmos3_bitwise_patches() -> None: + from miles.backends.sglang_diffusion_utils.monkey_patches import patch_cosmos3_bitwise + + patch_cosmos3_bitwise.apply() + + @register_rollout_patch_group("ltx") def apply_ltx2_rollout_patches() -> None: from miles.backends.sglang_diffusion_utils.monkey_patches import ( 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 new file mode 100644 index 00000000..9314e90c --- /dev/null +++ b/miles/backends/sglang_diffusion_utils/monkey_patches/patch_cosmos3_bitwise.py @@ -0,0 +1,192 @@ +"""Cosmos3 bitwise-parity patches for sgl-d (train-side reference: diffusers). + +Direction discipline (never downgrade precision): + +- Genuine precision-policy gaps are fixed on the LOW side. The only one found + is the train side's time_embedder (fixed there via the FSDP precision spec; + sgl-d already runs it fp32 — nothing to patch here). +- Kernel-organization differences are re-expressed on the sgl-d side as the + exact op sequence diffusers runs, at equal precision: + + * ``MergedColumnParallelLinear`` fuses Q/K/V (and gate/up) into one GEMM. + Measured on cosmos3's UND shapes (M=29, bf16, H200): the fused GEMM's Q + columns differ from the standalone Q GEMM by 3.5e-3 rel. Unfuse into + per-slice ``F.linear`` calls — each slice then runs the same GEMM the + diffusers module runs, bitwise. + * ``RMSNorm`` (flashinfer ``rmsnorm`` / ``fused_add_rmsnorm``) is rerouted + through ``F.rms_norm`` with the residual add made explicit in the input + dtype. The train side patches diffusers' RMSNorm onto the same + ``F.rms_norm`` (raising it from round-before-mul to fp32-through-mul), so + both stacks run the identical kernel. + * ``SiluAndMul`` (fused sgl-kernel, one rounding) becomes eager + ``F.silu(gate) * up`` (two roundings) — diffusers' exact op order. + * The fused qk-norm+rope JIT kernels are disabled; the split path runs the + same ``F.rms_norm`` + eager rope muls as diffusers. + * The GEN attention backend is pinned to TORCH_SDPA via backend selection + (the sanctioned channel — see monkey_patches.__init__ on why USPAttention + itself must not be patched); diffusers dispatches to the same + ``F.scaled_dot_product_attention``. + * CFG runs cond/uncond as two sequential batch-1 forwards instead of one + batch-2 forward. cuBLAS is not bitwise batch-invariant on cosmos3's + shapes (down_proj M=29->58 and proj_out M=390->780 both break), and the + 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. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + + +def apply() -> None: + _force_torch_sdpa_backend() + _patch_rmsnorm_f_rms_norm() + _patch_merged_column_linear_unfused() + _patch_silu_and_mul_eager() + _patch_qk_norm_rope_split_eager() + _patch_cfg_sequential() + + +def _force_torch_sdpa_backend() -> None: + from sglang.multimodal_gen.runtime.layers.attention.selector import global_force_attn_backend + from sglang.multimodal_gen.runtime.platforms.interface import AttentionBackendEnum + + global_force_attn_backend(AttentionBackendEnum.TORCH_SDPA) + + +def _patch_rmsnorm_f_rms_norm() -> None: + from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm + + def _forward(self, x: torch.Tensor, residual: torch.Tensor | None = None): + if self.variance_size_override is not None: + raise NotImplementedError("cosmos3_bitwise RMSNorm patch does not support variance_size_override") + if residual is not None: + # Same add the diffusers layer runs eagerly (single bf16 rounding). + residual = x + residual + out = F.rms_norm(residual, (self.hidden_size,), self.weight, self.variance_epsilon) + return out, residual + return F.rms_norm(x, (self.hidden_size,), self.weight, self.variance_epsilon) + + RMSNorm.forward_cuda = _forward + RMSNorm.forward_native = _forward + + +def _patch_merged_column_linear_unfused() -> None: + from sglang.multimodal_gen.runtime.layers.linear import MergedColumnParallelLinear, UnquantizedLinearMethod + + logged = False + + def _forward(self, x: torch.Tensor): + # NOT self.output_partition_sizes: with tp=1, MergedColumnParallelLinear + # assigns self.output_sizes only after super().__init__() has already + # derived output_partition_sizes, so that attr collapses to + # [sum(output_sizes)] and a "per-slice" loop over it degenerates into + # the very fused GEMM this patch exists to avoid. + sizes = getattr(self, "output_sizes", None) + if not isinstance(self.quant_method, UnquantizedLinearMethod) or self.skip_bias_add or sizes is None: + raise RuntimeError( + "cosmos3_bitwise unfused-GEMM patch cannot handle this MergedColumnParallelLinear " + f"(quant_method={type(self.quant_method).__name__}, skip_bias_add={self.skip_bias_add}, " + f"output_sizes={sizes}); refusing to fall back to the fused GEMM silently." + ) + sizes = [size // self.tp_size for size in sizes] + nonlocal logged + if not logged: + logged = True + print(f"[cosmos3_bitwise] unfused MergedColumnParallelLinear active: slices={sizes}", flush=True) + outs = [] + offset = 0 + for size in 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)) + offset += size + return torch.cat(outs, dim=-1), None + + MergedColumnParallelLinear.forward = _forward + + +def _patch_silu_and_mul_eager() -> None: + from sglang.multimodal_gen.runtime.layers.activation import SiluAndMul + + def _forward(self, x: torch.Tensor) -> torch.Tensor: + d = x.shape[-1] // 2 + return F.silu(x[..., :d]) * x[..., d:] + + SiluAndMul.forward_cuda = _forward + SiluAndMul.forward_native = _forward + + +def _patch_qk_norm_rope_split_eager() -> None: + from sglang.multimodal_gen.runtime.layers import layernorm + from sglang.multimodal_gen.runtime.models.dits import cosmos3video + + # The split path's apply_qk_norm falls back to the (patched) RMSNorm + # modules once the fused inplace JIT kernel is declared unavailable. + layernorm.can_use_fused_inplace_qknorm = lambda *args, **kwargs: False + + def _delegate(q, k, q_norm, k_norm, head_dim, cos_sin_cache, rope_cache_positions): + return cosmos3video._apply_qwen3_qk_norm_rope_split(q, k, q_norm, k_norm, head_dim, cos_sin_cache) + + cosmos3video._apply_qwen3_qk_norm_rope = _delegate + + +def _dumper_step_between_branches() -> None: + # The denoising stage's Dumper instrumentation steps once per loop + # iteration; sequential CFG puts two forwards in one iteration, which + # would collide record names. Step between the branches so every dumper + # step holds exactly one forward (uncond and cond land in adjacent steps; + # the comparator pairs by bit-exact latent anchor, not by step index). + try: + from sglang.srt.debug_utils.dumper import dumper + except ImportError: + return + if dumper.may_enable and dumper._non_intrusives: + dumper.step() + + +def _patch_cfg_sequential() -> None: + from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.cosmos3 import ( + Cosmos3DenoisingStage, + ) + + def _predict_noise_cfg_batched( + self, + latents: torch.Tensor, + timestep: torch.Tensor, + cond_text_ids: torch.Tensor, + cond_text_mask: torch.Tensor, + uncond_text_ids: torch.Tensor, + uncond_text_mask: torch.Tensor, + video_shape: tuple[int, int, int], + fps: float, + guidance_scale: float, + noisy_frame_mask: torch.Tensor | None = None, + max_text_seq_len: int | None = None, + current_timestep: int | None = None, + ) -> torch.Tensor: + del max_text_seq_len # per-branch true length, recomputed from each mask + + def run(text_ids, text_mask, cache_key): + return self._run_transformer( + latents=latents, + timestep=timestep, + text_ids=text_ids, + text_mask=text_mask, + video_shape=video_shape, + fps=fps, + cache_key=cache_key, + noisy_frame_mask=noisy_frame_mask, + max_text_seq_len=None, + current_timestep=current_timestep, + ) + + noise_pred_uncond = run(uncond_text_ids, uncond_text_mask, "uncond") + _dumper_step_between_branches() + noise_pred_cond = run(cond_text_ids, cond_text_mask, "cond") + # CFG: uncond + g·(cond − uncond) — same op order as the train side's + # cfg_combine and the original batched combine. + return noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond) + + Cosmos3DenoisingStage._predict_noise_cfg_batched = _predict_noise_cfg_batched diff --git a/scripts/run-diffusion-grpo-cosmos3-pickscore-t2i-5gpu.sh b/scripts/run-diffusion-grpo-cosmos3-pickscore-t2i-5gpu.sh index dc796fd1..a8ffbc8e 100755 --- a/scripts/run-diffusion-grpo-cosmos3-pickscore-t2i-5gpu.sh +++ b/scripts/run-diffusion-grpo-cosmos3-pickscore-t2i-5gpu.sh @@ -93,6 +93,7 @@ fi --use-miles-router \ --sglang-server-concurrency 8 \ --sglang-attention-backend fa \ + --rollout-patch-group cosmos3_bitwise \ --update-weight-buffer-size 2147483648 \ --update-weight-target-module transformer \ --diffusion-reward pickscore:1.0 \ From 6d60e3d9a98f25744756183cdf2c2e897aa69229 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Tue, 11 Aug 2026 02:46:35 +0000 Subject: [PATCH 2/3] feat(cosmos3): bitwise parity across LoRA updates via the native LoRA-IPC path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit GEMM(W + s·BA) is not bitwise GEMM(W) + GEMM_B(GEMM_A(x))·s, so the default lora_merge weight sync caps train/rollout parity at the first step (lora_B starts at zero). Keep the engine-side adapters unmerged instead, mirroring the qwen_image patch group (#108): - recipes ship adapters with --lora-ipc-weight-sync — fp32 lora_A/lora_B masters through the engine's native LoRA-IPC path; no bespoke transport. - cosmos3_bitwise patches the native LoRA wrappers: set_lora_weights never merges and rounds A/B to the base weight dtype (the FSDP mixed-precision gather rounding the train forward sees); wrapper forwards run eager base(x) + lora_B(lora_A(x))·s in peft's exact op order (the stock forwards are @torch.compile'd, which re-fuses even the no-adapter base path). - fused targets (add_q/k/v -> to_qkv) all resolve to one wrapper whose set_lora_weights(clear_existing=True) calls would clobber each other; _resolve_lora_ipc_layer_dict_key is patched to route each prefix to its merge slot, and the delta lands on the matching output slice. - adapt the CFG-sequential patch to the omni-era _run_transformer kwargs and guard the new fused qknorm+rope kernel — current sglang main moved both. - fix the 5gpu recipe's stale --diffusion-init-lora-weight flag (renamed to --lora-init-weights on main). Verified on Cosmos3-Nano GRPO (3 train GPUs + pickscore): LoRA IPC sync resolves all 144 layer prefixes (unmapped=0) and train/model_output_{mean,max}_abs_diff stay 0.0 across steps 1-3, i.e. across two LoRA weight updates. Requires engine-side Cosmos3Pipeline LoRA support (the LoRAPipeline mixin), added to sgl-project/sglang#34197. Supersedes the bespoke unmerged-sync transport from #129. Co-authored-by: Cursor --- .../monkey_patches/patch_cosmos3_bitwise.py | 177 +++++++++++++++++- ...ffusion-grpo-cosmos3-pickscore-t2i-5gpu.sh | 3 +- 2 files changed, 177 insertions(+), 3 deletions(-) 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..e12b5ec5 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,17 @@ 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-ipc-weight-sync`` the trainer + ships the fp32 lora_A/lora_B masters through the engine's native LoRA-IPC + path (weights_updater -> LoRAPipeline wrappers, same as qwen_image). This + side keeps the wrappers unmerged, rounds A/B to the base weight dtype (the + rounding FSDP's mixed-precision gather applies before the train forward), + and replays peft's exact op sequence per target — fused targets + (add_q/k/v -> to_qkv) are routed to their merge slot and added on the + matching output slice. """ from __future__ import annotations @@ -47,6 +58,7 @@ def apply() -> None: _patch_silu_and_mul_eager() _patch_qk_norm_rope_split_eager() _patch_cfg_sequential() + _patch_lora_unmerged_native() def _force_torch_sdpa_backend() -> None: @@ -107,6 +119,149 @@ def _forward(self, x: torch.Tensor): 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's forward sees the fp32 adapter masters FSDP-gathered at the + forward dtype; the wrappers store A/B rounded to that same dtype (see + ``_round_adapter``), 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 _round_adapter(t: torch.Tensor, layer) -> torch.Tensor: + # fp32 master -> the same elementwise cast FSDP's mixed-precision gather + # applies before the train forward. copy=True also detaches the tensor + # from the CUDA-IPC flattened bucket, whose storage the sender reclaims + # after the update returns (the updater ipc_collects imported blocks). + return t.to(dtype=layer.base_layer.weight.dtype, copy=True).contiguous() + + +def _wrapper_scaling(layer) -> float: + # peft's ``scaling = lora_alpha / r`` (exact python-float division); the + # engine-side strength knob stays folded in for completeness (1.0 here). + scaling = layer.lora_alpha / layer.lora_rank + if layer.strength != 1.0: + scaling = scaling * layer.strength + return scaling + + +class _FusedSlotView: + """Stand-in for one merge slot of a fused LoRA wrapper. + + The LoRA-IPC update loop resolves add_q/k/v_proj to the same to_qkv + wrapper and calls ``set_lora_weights(..., clear_existing=True)`` — the + slots would clobber each other. ``_resolve_lora_ipc_layer_dict_key`` is + patched to hand the loop this view instead, routing each adapter to its + slot on the real wrapper. The loop assigns ``lora_rank`` / ``lora_alpha`` + on the view before calling set_lora_weights; plain attribute writes keep + them off the shared wrapper. + """ + + def __init__(self, layer, slot: int): + self._layer = layer + self._slot = slot + + def set_lora_weights(self, A, B, lora_path=None, strength=1.0, clear_existing=False, merge_weights=True): + del lora_path, clear_existing, merge_weights + layer = self._layer + scaling = self.lora_alpha / self.lora_rank + if strength != 1.0: + scaling = scaling * strength + slots = getattr(layer, "_miles_lora_slots", None) + if slots is None: + slots = {} + layer._miles_lora_slots = slots + slots[self._slot] = (_round_adapter(A, layer), _round_adapter(B, layer), scaling) + + +def _patch_lora_unmerged_native() -> None: + """Keep engine-side LoRA unmerged and bitwise-equal to the peft train path. + + The trainer's ``--lora-ipc-weight-sync`` ships the fp32 lora_A/lora_B + masters through the engine's native LoRA-IPC path (weights_updater -> + LoRAPipeline/BaseLayerWithLoRA — the same mechanism the qwen_image group + patches). Stock behaviour merges ``W' = W + s·BA`` in bf16, which caps + train/rollout parity at the first step: ``GEMM(W + sBA)`` is not bitwise + ``GEMM(W) + GEMM_B(GEMM_A(x))·s``. Mirror the train-side compute instead: + + - ``set_lora_weights``: never merge; round A/B to the base weight dtype + (the rounding FSDP's mixed-precision gather applies on the train side). + - forwards: eager ``base(x) + lora_B(lora_A(x))·s`` in peft's op order. + The stock wrappers run under ``@torch.compile``, which re-fuses even the + no-adapter base path — every wrapper forward must be replaced. + - fused targets (add_q/k/v -> to_qkv): route each prefix to its merge slot + via ``_FusedSlotView`` and add its delta on the matching output slice — + elementwise identical to the train side's per-projection ``base + delta`` + before concat. + + TP=1 only, like the rest of this patch group (adapters are stored and + applied unsliced). + """ + from sglang.multimodal_gen.runtime.layers.lora.linear import ( + BaseLayerWithLoRA, + ColumnParallelLinearWithLoRA, + LinearWithLoRA, + MergedColumnParallelLinearWithLoRA, + RowParallelLinearWithLoRA, + ) + from sglang.multimodal_gen.runtime.post_training import weights_updater + + orig_set_lora_weights = BaseLayerWithLoRA.set_lora_weights + + def _set_lora_weights_unmerged(self, A, B, *args, **kwargs): + kwargs["merge_weights"] = False + return orig_set_lora_weights(self, _round_adapter(A, self), _round_adapter(B, self), *args, **kwargs) + + BaseLayerWithLoRA.set_lora_weights = _set_lora_weights_unmerged + + def _tuple_lora_forward(self, x: torch.Tensor): + out, output_bias = self.base_layer(x) + if not self.merged and not self.disable_lora: + # After the complete base output (bias included) — the position + # peft adds the delta at. + out = out + _lora_term(x, self.lora_A, self.lora_B, _wrapper_scaling(self)) + return out, output_bias + + def _nn_linear_lora_forward(self, x: torch.Tensor): + out = self.base_layer(x) + if not self.merged and not self.disable_lora: + out = out + _lora_term(x, self.lora_A, self.lora_B, _wrapper_scaling(self)) + return out + + def _merged_column_lora_forward(self, x: torch.Tensor): + out, output_bias = self.base_layer(x) + slots = getattr(self, "_miles_lora_slots", None) + if slots: + sizes = [size // self.base_layer.tp_size for size in self.base_layer.output_sizes] + for slot, (A, B, scaling) in sorted(slots.items()): + offset = sum(sizes[:slot]) + out[..., offset : offset + sizes[slot]] += _lora_term(x, A, B, scaling) + return out, output_bias + + BaseLayerWithLoRA.forward = _tuple_lora_forward + ColumnParallelLinearWithLoRA.forward = _tuple_lora_forward + RowParallelLinearWithLoRA.forward = _tuple_lora_forward + MergedColumnParallelLinearWithLoRA.forward = _merged_column_lora_forward + LinearWithLoRA.forward = _nn_linear_lora_forward + + orig_resolve = weights_updater._resolve_lora_ipc_layer_dict_key + + def _resolve_with_slots(layer_prefix, layer_dict, module): + layer, key = orig_resolve(layer_prefix, layer_dict, module) + if layer is None: + return layer, key + map_name = weights_updater._build_module_weight_name_mapper(module) + slot = map_name(f"{layer_prefix}.weight")[1] if map_name is not None else None + if slot is None: + return layer, key + return _FusedSlotView(layer, slot), key + + weights_updater._resolve_lora_ipc_layer_dict_key = _resolve_with_slots + + def _patch_silu_and_mul_eager() -> None: from sglang.multimodal_gen.runtime.layers.activation import SiluAndMul @@ -125,6 +280,11 @@ def _patch_qk_norm_rope_split_eager() -> None: # The split path's apply_qk_norm falls back to the (patched) RMSNorm # modules once the fused inplace JIT kernel is declared unavailable. layernorm.can_use_fused_inplace_qknorm = lambda *args, **kwargs: False + if hasattr(layernorm, "can_use_fused_inplace_qknorm_rope"): + # Newer trees add a fused qknorm+rope kernel with its own guard; the + # delegate below bypasses its only cosmos3 call site, this is belt and + # braces should another path reach apply_qk_norm_rope. + layernorm.can_use_fused_inplace_qknorm_rope = lambda *args, **kwargs: False def _delegate(q, k, q_norm, k_norm, head_dim, cos_sin_cache, rope_cache_positions): return cosmos3video._apply_qwen3_qk_norm_rope_split(q, k, q_norm, k_norm, head_dim, cos_sin_cache) @@ -165,8 +325,14 @@ def _predict_noise_cfg_batched( noisy_frame_mask: torch.Tensor | None = None, max_text_seq_len: int | None = None, current_timestep: int | None = None, - ) -> torch.Tensor: + **extra, + ) -> torch.Tensor | tuple[torch.Tensor, ...]: del max_text_seq_len # per-branch true length, recomputed from each mask + # Omni-era conditioning (sound/action latents and friends) is identical + # across CFG branches — the batched impl torch.cat's each with itself — + # so it passes through per-branch unchanged. Drop the Nones so the same + # code runs on trees whose _run_transformer predates these kwargs. + extra = {key: value for key, value in extra.items() if value is not None} def run(text_ids, text_mask, cache_key): return self._run_transformer( @@ -180,13 +346,20 @@ def run(text_ids, text_mask, cache_key): noisy_frame_mask=noisy_frame_mask, max_text_seq_len=None, current_timestep=current_timestep, + **extra, ) noise_pred_uncond = run(uncond_text_ids, uncond_text_mask, "uncond") _dumper_step_between_branches() noise_pred_cond = run(cond_text_ids, cond_text_mask, "cond") + # CFG: uncond + g·(cond − uncond) — same op order as the train side's # cfg_combine and the original batched combine. - return noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond) + def combine(uncond, cond): + return uncond + guidance_scale * (cond - uncond) + + if isinstance(noise_pred_cond, tuple): + return tuple(combine(u, c) for u, c in zip(noise_pred_uncond, noise_pred_cond, strict=True)) + return combine(noise_pred_uncond, noise_pred_cond) Cosmos3DenoisingStage._predict_noise_cfg_batched = _predict_noise_cfg_batched diff --git a/scripts/run-diffusion-grpo-cosmos3-pickscore-t2i-5gpu.sh b/scripts/run-diffusion-grpo-cosmos3-pickscore-t2i-5gpu.sh index a8ffbc8e..3103e83a 100755 --- a/scripts/run-diffusion-grpo-cosmos3-pickscore-t2i-5gpu.sh +++ b/scripts/run-diffusion-grpo-cosmos3-pickscore-t2i-5gpu.sh @@ -84,7 +84,8 @@ fi --use-lora \ --lora-rank 64 \ --lora-alpha 128 \ - --diffusion-init-lora-weight gaussian \ + --lora-ipc-weight-sync \ + --lora-init-weights gaussian \ --lr 3e-4 \ --adam-beta2 0.95 \ --clip-grad 2e-3 \ From 8a9e638235950aee26fff8651b5452f9b3a3539e Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Wed, 12 Aug 2026 07:03:44 +0000 Subject: [PATCH 3/3] fix pre-commit formatting --- miles/backends/fsdp_utils/actor.py | 4 +++- .../monkey_patches/patch_cosmos3_bitwise.py | 4 +--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index f012bbbf..b75392a2 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -98,7 +98,9 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty self._master_dtype = parse_dtype_from_str(args.fsdp_master_dtype) self._forward_dtype = parse_dtype_from_str(args.diffusion_forward_dtype) - self._frozen_dtype = parse_dtype_from_str(args.fsdp_frozen_params_dtype) if args.fsdp_frozen_params_dtype else None + self._frozen_dtype = ( + parse_dtype_from_str(args.fsdp_frozen_params_dtype) if args.fsdp_frozen_params_dtype else None + ) if self._frozen_dtype is not None and not args.use_lora: raise ValueError("--fsdp-frozen-params-dtype requires --use-lora") 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 e12b5ec5..08542b5d 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 @@ -307,9 +307,7 @@ def _dumper_step_between_branches() -> None: def _patch_cfg_sequential() -> None: - from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.cosmos3 import ( - Cosmos3DenoisingStage, - ) + from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.cosmos3 import Cosmos3DenoisingStage def _predict_noise_cfg_batched( self,