diff --git a/miles/backends/fsdp_utils/configs/cosmos3.py b/miles/backends/fsdp_utils/configs/cosmos3.py index 49a4f442..97d6f304 100644 --- a/miles/backends/fsdp_utils/configs/cosmos3.py +++ b/miles/backends/fsdp_utils/configs/cosmos3.py @@ -5,6 +5,7 @@ import math import torch +from miles.backends.sglang_diffusion_utils.monkey_patches import patch_cosmos3_bitwise from miles.utils.types import CondKwargs from .train_pipeline_config import TrainPipelineConfig, register_train_pipeline_config @@ -49,6 +50,9 @@ def validate_args(cls, args) -> None: if list(args.update_weight_target_modules) != ["transformer"]: raise ValueError("Cosmos3 requires --update-weight-target-module transformer.") + def configure(self, args) -> None: + self._bitwise_parity = "cosmos3_bitwise" in args.rollout_patch_groups + def prepare_cond_kwargs(self, cond: CondKwargs | None, device: torch.device) -> dict: if cond is None or cond.text_ids is None: return {} @@ -175,3 +179,5 @@ def _cast_to_weight_dtype(module, args): 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) + if self._bitwise_parity: + patch_cosmos3_bitwise.apply_train(model) 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..3513da82 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,6 @@ from miles.backends.fsdp_utils.models.parallel_plan import FSDPParallelPlan -FSDP_PARALLEL_PLAN = FSDPParallelPlan() +FSDP_PARALLEL_PLAN = FSDPParallelPlan( + param_dtype_patterns={"*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..4b317b39 --- /dev/null +++ b/miles/backends/sglang_diffusion_utils/monkey_patches/patch_cosmos3_bitwise.py @@ -0,0 +1,402 @@ +"""Cosmos3 bitwise-parity patches: the sgl-d rollout group (`apply`) and the train-side halves (`apply_train`). + +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. + * 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 + +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() + _patch_lora_unmerged_native() + + +def apply_train(model: torch.nn.Module) -> None: + """Train-side halves of the parity contract; called from the cosmos3 config's postprocess.""" + _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: + """Dedupe identical per-token sinusoid rows to sglang-d's M=1 GEMM shape (cuBLAS is not M-invariant).""" + orig_forward = time_embedder.forward + + def forward(x, *args, **kwargs): + # Autocast off: bf16 autocast would undo the fp32 weight gather at the matmul boundary. + 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: + """Route diffusers RMSNorm through F.rms_norm (fp32 through the weight mul), same op as _patch_rmsnorm_f_rms_norm.""" + 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 + + +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 _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 + + 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 + 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) + + 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, + **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( + 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, + **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. + 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 375f0fcd..9521de35 100755 --- a/scripts/run-diffusion-grpo-cosmos3-pickscore-t2i-5gpu.sh +++ b/scripts/run-diffusion-grpo-cosmos3-pickscore-t2i-5gpu.sh @@ -83,7 +83,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 \ @@ -92,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 \