From 46f15e42a7c6c3d09c38457329dd41e834f2be12 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Mon, 17 Aug 2026 22:33:43 +0000 Subject: [PATCH 1/5] feat(cosmos3): bitwise train/rollout parity via the cosmos3_bitwise rollout patch group --- miles/backends/fsdp_utils/configs/cosmos3.py | 6 + .../models/diffusers/cosmos3/parallel_plan.py | 4 +- .../monkey_patches/__init__.py | 7 + .../monkey_patches/patch_cosmos3_bitwise.py | 372 ++++++++++++++++++ miles/utils/arguments.py | 8 + .../test_rollout_patch_groups.py | 1 + 6 files changed, 397 insertions(+), 1 deletion(-) 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 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 edbc3c97..e69ab4f1 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("wan") def apply_wan_rollout_patches() -> None: from miles.backends.sglang_diffusion_utils.monkey_patches import patch_wan_norm_ops 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..d5eed971 --- /dev/null +++ b/miles/backends/sglang_diffusion_utils/monkey_patches/patch_cosmos3_bitwise.py @@ -0,0 +1,372 @@ +"""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). The recipe keeps the wrappers unmerged via + ``--sglang-lora-merge-mode dynamic``; with ``--lora-ipc-weight-sync`` the + trainer ships lora_A/lora_B through the engine's native LoRA-IPC path, + rounded to the train forward dtype before send (the rounding FSDP's + mixed-precision gather applies before the train forward). The wrapper + forwards replay peft's exact op sequence per target — fused targets + (add_q/k/v -> to_qkv) arrive as one block-diagonal composed pair and each + section's delta lands on its 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_peft_forwards() + + +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() + _round_lora_ipc_sends_to_forward_dtype() + + +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 sender ships A/B rounded to that same dtype (see + ``_round_lora_ipc_sends_to_forward_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 _round_lora_ipc_sends_to_forward_dtype() -> None: + from miles.backends.fsdp_utils.diffusion_update_weight_utils import DiffusionUpdateWeightFromTensorLoRAIPC + from miles.backends.fsdp_utils.mixed_precision import parse_dtype_from_str + + if getattr(DiffusionUpdateWeightFromTensorLoRAIPC, "_miles_rounded_lora_send", False): + return + + orig_prepare = DiffusionUpdateWeightFromTensorLoRAIPC._prepare_lora_param + + def _prepare(self, param: torch.Tensor) -> torch.Tensor: + return orig_prepare(self, param).to(parse_dtype_from_str(self.args.diffusion_forward_dtype)) + + DiffusionUpdateWeightFromTensorLoRAIPC._prepare_lora_param = _prepare + DiffusionUpdateWeightFromTensorLoRAIPC._miles_rounded_lora_send = True + + +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 + + +def _patch_lora_peft_forwards() -> None: + """Replay peft's unmerged op order in the engine's LoRA wrappers. + + Stock wrapper forwards 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) arrive as one block-diagonal composed pair + (scale folded into B, exact for power-of-two alpha/rank); each section's + delta lands on its output slice — elementwise identical to the train + side's per-projection ``base + delta`` before concat. Bitwise parity is a + tp_size==1 property; sharded layers fall back to the native TP-aware + forwards. + """ + from sglang.multimodal_gen.runtime.layers.lora.linear import ( + BaseLayerWithLoRA, + ColumnParallelLinearWithLoRA, + LinearWithLoRA, + MergedColumnParallelLinearWithLoRA, + RowParallelLinearWithLoRA, + ) + + orig_column_forward = ColumnParallelLinearWithLoRA.forward + orig_row_forward = RowParallelLinearWithLoRA.forward + orig_merged_forward = MergedColumnParallelLinearWithLoRA.forward + + 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 _column_parallel_lora_forward(self, x: torch.Tensor): + if self.base_layer.tp_size > 1: + return orig_column_forward(self, x) + return _tuple_lora_forward(self, x) + + def _row_parallel_lora_forward(self, x: torch.Tensor): + if self.base_layer.tp_size > 1: + return orig_row_forward(self, x) + return _tuple_lora_forward(self, x) + + def _merged_column_lora_forward(self, x: torch.Tensor): + if self.base_layer.tp_size > 1: + return orig_merged_forward(self, x) + out, output_bias = self.base_layer(x) + if not self.merged and not self.disable_lora: + sizes = self.base_layer.output_sizes + rank = self.lora_A.shape[0] // len(sizes) + scaling = _wrapper_scaling(self) + row = col = 0 + for size in sizes: + a = self.lora_A[col : col + rank] + # Contiguous copy keeps the GEMM layout identical to the train + # side's standalone per-projection GEMM. + b = self.lora_B[row : row + size, col : col + rank].contiguous() + out[..., row : row + size] += _lora_term(x, a, b, scaling) + row += size + col += rank + return out, output_bias + + BaseLayerWithLoRA.forward = _tuple_lora_forward + ColumnParallelLinearWithLoRA.forward = _column_parallel_lora_forward + RowParallelLinearWithLoRA.forward = _row_parallel_lora_forward + MergedColumnParallelLinearWithLoRA.forward = _merged_column_lora_forward + LinearWithLoRA.forward = _nn_linear_lora_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 + 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/miles/utils/arguments.py b/miles/utils/arguments.py index deeb75f9..b2e64a15 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1509,6 +1509,14 @@ def miles_validate_args(args): from miles.backends.sglang_diffusion_utils.monkey_patches import validate_rollout_patch_groups validate_rollout_patch_groups(args.rollout_patch_groups) + if args.use_lora and "cosmos3_bitwise" in args.rollout_patch_groups: + # Missing on engines whose ServerArgs predates --lora-merge-mode. + if getattr(args, "sglang_lora_merge_mode", None) != "dynamic": + logger.warning( + "cosmos3_bitwise runs LoRA without --sglang-lora-merge-mode dynamic; the engine " + "auto-merges the adapters into the base weights, introducing precision " + "drift against the trainer's unmerged forward — training still works." + ) if args.lora_ipc_weight_sync: if not args.use_lora: diff --git a/tests/fast/backends/sglang_diffusion_utils/test_rollout_patch_groups.py b/tests/fast/backends/sglang_diffusion_utils/test_rollout_patch_groups.py index f844fdf6..09f31e7d 100644 --- a/tests/fast/backends/sglang_diffusion_utils/test_rollout_patch_groups.py +++ b/tests/fast/backends/sglang_diffusion_utils/test_rollout_patch_groups.py @@ -44,6 +44,7 @@ def test_unknown_group_fails_loud(self, monkeypatch): def test_builtin_group_registered(self): # The decorator ran at import time for the in-repo groups. assert "sgld" in mp._ROLLOUT_PATCH_APPLIERS + assert "cosmos3_bitwise" in mp._ROLLOUT_PATCH_APPLIERS assert "wan" in mp._ROLLOUT_PATCH_APPLIERS From bffbb7e0128cfdf5ef2886068887b2719bba6789 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Mon, 17 Aug 2026 22:33:43 +0000 Subject: [PATCH 2/5] add the cosmos3 bitwise 4gpu T2I recipe script --- ...grpo_cosmos3_pickscore_t2i_4gpu_bitwise.py | 168 ++++++++++++++++++ 1 file changed, 168 insertions(+) create mode 100644 scripts/run_diffusion_grpo_cosmos3_pickscore_t2i_4gpu_bitwise.py diff --git a/scripts/run_diffusion_grpo_cosmos3_pickscore_t2i_4gpu_bitwise.py b/scripts/run_diffusion_grpo_cosmos3_pickscore_t2i_4gpu_bitwise.py new file mode 100644 index 00000000..55640be4 --- /dev/null +++ b/scripts/run_diffusion_grpo_cosmos3_pickscore_t2i_4gpu_bitwise.py @@ -0,0 +1,168 @@ +"""Cosmos3-Nano T2I GRPO with PickScore, fully colocated on 4 GPUs, with +bitwise train<->rollout parity. + +pretrained = nvidia/Cosmos3-Nano (16B MoT: 8B UND tower frozen, 8B GEN tower +trained via LoRA r64), 832x480 single frame, num_steps=16, eval_steps=35, +guidance 4.0, Flow-SDE noise_level=0.7, no KL, per-prompt mean + global std. + +Layout: train, rollout and PickScore reward all share the same 4 GPUs +(--colocate --colocate-reward, one PickScore worker per rollout engine). + +SDE schedule: epoch_global_random_choice draws 2 steps per epoch from +candidates 4-7 — the high-noise segment (sigma 0.94-0.80) of the FlowUniPC +shift-3 grid the rollout inherits from serving. Step numbers are NOT +transferable across sigma-grid families - re-derive candidates from |dt| when +changing model/grid. + +Pacing: lr 1e-4 x 1 optimizer step per rollout (the whole rollout is one +batch). CFG amplifies per-step policy displacement, so training with +guidance > 1 needs this slower pacing than a comparable CFG-free recipe. + +Bitwise parity: --rollout-patch-group cosmos3_bitwise re-expresses the +engine's kernel organization as the exact op sequence the trainer runs, and +--sglang-lora-merge-mode dynamic keeps the adapters unmerged (merging rounds +differently in bf16). Parity holds at tp_size==1 with --lora-ipc-weight-sync. + +--diffusion-recompute-old-log-prob: the trainer recomputes old log-probs at +rollout ingestion so the PPO ratio is implementation-self-consistent; with +parity on both sides agree bitwise, so this only guards runs without the +patch group. With 1 step per rollout every optimizer step is exactly +on-policy. + +Usage: + python3 scripts/run_diffusion_grpo_cosmos3_pickscore_t2i_4gpu_bitwise.py +""" + +from dataclasses import dataclass + +import typer + +import miles.utils.external_utils.command_utils as U + +MODEL = "nvidia/Cosmos3-Nano" +DATASET = "rockdu/miles-diffusion-datasets" +DATASET_SUBSET = "flowgrpo_pickscore" +WANDB_PROJECT = "miles-diffusion-grpo" + + +@dataclass +class ScriptArgs(U.ExecuteTrainConfig): + cuda_visible_devices: str = "0,1,2,3" + num_rollout: int = 10000 + data_dir: str = "/root/datasets" + extra_args: str = "" + + +def prepare(args: ScriptArgs) -> str: + local_dir = U.hf_download_dataset(DATASET, include=f"{DATASET_SUBSET}/**", data_dir=args.data_dir) + return f"{local_dir}/{DATASET_SUBSET}" + + +def execute(args: ScriptArgs, data_dir: str) -> None: + run_name = f"diffusion_grpo_cosmos3_pickscore_t2i_4gpu_bitwise_{U.create_run_id()}" + + ckpt_args = f"--hf-checkpoint {MODEL} --save {args.output_dir}/{run_name}/ckpt --save-interval 10 " + + rollout_args = ( + "--rollout-function-path miles.rollout.sglang_diffusion_rollout.generate_rollout " + f"--prompt-data {data_dir}/train.jsonl " + "--input-key input " + "--rollout-batch-size 48 " + "--n-samples-per-prompt 16 " + f"--num-rollout {args.num_rollout} " + "--num-steps-per-rollout 1 " + # The Cosmos3 transformer is a packed-sequence single-sample interface; + # one request cannot batch multiple outputs. + "--rollout-microgroup-size 1 " + "--micro-batch-size 1 " + ) + + diffusion_args = ( + "--diffusion-num-steps 16 " + "--diffusion-output-num-frames 1 " + "--diffusion-guidance-scale 4.0 " + "--diffusion-noise-level 0.7 " + "--diffusion-height 480 " + "--diffusion-width 832 " + "--diffusion-step-strategy-path miles.rollout.step_strategy_hub.epoch_global_random_choice " + "--diffusion-num-sde-steps 2 " + "--diffusion-sde-candidate-steps 4,5,6,7 " + "--diffusion-recompute-old-log-prob " + ) + + parity_args = "--rollout-patch-group cosmos3_bitwise --sglang-lora-merge-mode dynamic " + + eval_args = ( + f"--eval-prompt-data pickscore_test {data_dir}/test.jsonl " + "--eval-interval 30 " + "--diffusion-eval-num-steps 35 " + "--skip-eval-before-train " + ) + + grpo_args = "--advantage-estimator grpo --globalize-reward-std --diffusion-clip-range 1e-3 " + + optimizer_args = "--lr 1e-4 --adam-beta2 0.999 --weight-decay 1e-4 " + + # UND/GEN towers share layers and differ by parameter name (to_q vs + # add_q_proj, mlp vs mlp_moe_gen); LoRA targeting defaults to the GEN + # fragments in the cosmos3 train pipeline config. + lora_args = "--use-lora --lora-ipc-weight-sync --lora-rank 64 --lora-alpha 128 --lora-init-weights gaussian " + + reward_args = ( + "--rm-type pickscore " + "--colocate-reward " + "--pickscore-num-workers 4 " + "--pickscore-batch-size 8 " + "--pickscore-processor-path laion/CLIP-ViT-H-14-laion2B-s32B-b79K " + "--pickscore-model-path yuvalkirstain/PickScore_v1 " + ) + + wandb_args = U.get_default_wandb_args( + __file__, run_id=run_name, project=WANDB_PROJECT, wandb_log_num_images=8, wandb_log_image_interval=10 + ) + + sglang_args = ( + "--use-miles-router " + "--sglang-server-concurrency 8 " + "--update-weight-buffer-size 2147483648 " + "--update-weight-target-module transformer " + ) + + train_backend_args = ( + "--train-backend fsdp --fsdp-master-dtype fp32 --fsdp-reduce-dtype fp32 --diffusion-forward-dtype bf16 " + ) + + misc_args = ( + "--actor-num-gpus-per-node 4 " + "--rollout-num-gpus 4 " + "--rollout-num-gpus-per-engine 1 " + "--num-gpus-per-node 4 " + "--colocate " + ) + + debug_args = "--diffusion-debug-mode " + + U.execute_train( + train_args=( + f"{ckpt_args} {rollout_args} {diffusion_args} {parity_args} {eval_args} {grpo_args} " + f"{optimizer_args} {lora_args} {reward_args} {wandb_args} {sglang_args} {train_backend_args} " + f"{misc_args} {debug_args} {args.extra_args}" + ), + num_gpus_per_node=4, + config=args, + extra_env_vars={ + "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:False", + # RL rollout scores raw samples; skip the serving-side guardrail models. + "SGLANG_DISABLE_COSMOS3_GUARDRAILS": "1", + }, + ) + + +@U.dataclass_cli +def main(args: ScriptArgs) -> None: + data_dir = prepare(args) + execute(args, data_dir) + + +if __name__ == "__main__": + typer.run(main) From 71cf0bd4db994f399c8dfd590eb8b5f1dab1a680 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Tue, 18 Aug 2026 09:21:15 +0000 Subject: [PATCH 3/5] ci: add the cosmos3 t2i pickscore GRPO e2e test Registers the bitwise 4-GPU recipe in stage-c-5-gpu-h200 (opt-in via run-ci-e2e, nightly). --deterministic-mode joins the recipe so the standard compares strictly; the model_output parity metrics must be exactly 0 under the cosmos3_bitwise patch group + dynamic LoRA merge. Co-Authored-By: Claude Fable 5 --- ...grpo_cosmos3_pickscore_t2i_4gpu_bitwise.py | 2 +- .../test_cosmos3_pickscore_grpo_t2i_4xGPU.py | 40 +++++++++++++++++++ 2 files changed, 41 insertions(+), 1 deletion(-) create mode 100644 tests/e2e/short/test_cosmos3_pickscore_grpo_t2i_4xGPU.py diff --git a/scripts/run_diffusion_grpo_cosmos3_pickscore_t2i_4gpu_bitwise.py b/scripts/run_diffusion_grpo_cosmos3_pickscore_t2i_4gpu_bitwise.py index 55640be4..b1e8929f 100644 --- a/scripts/run_diffusion_grpo_cosmos3_pickscore_t2i_4gpu_bitwise.py +++ b/scripts/run_diffusion_grpo_cosmos3_pickscore_t2i_4gpu_bitwise.py @@ -140,7 +140,7 @@ def execute(args: ScriptArgs, data_dir: str) -> None: "--colocate " ) - debug_args = "--diffusion-debug-mode " + debug_args = "--deterministic-mode --diffusion-debug-mode " U.execute_train( train_args=( diff --git a/tests/e2e/short/test_cosmos3_pickscore_grpo_t2i_4xGPU.py b/tests/e2e/short/test_cosmos3_pickscore_grpo_t2i_4xGPU.py new file mode 100644 index 00000000..848e7f56 --- /dev/null +++ b/tests/e2e/short/test_cosmos3_pickscore_grpo_t2i_4xGPU.py @@ -0,0 +1,40 @@ +"""E2E: Cosmos3-Nano T2I PickScore GRPO, 4-GPU fully colocated (train, rollout +and PickScore reward share the same 4 GPUs) — runs the bitwise recipe's real +configuration and checks its metric series against the registered standard +(tests/ci/fixtures/e2e_standards/). Runs with --deterministic-mode, so every +metric is compared strictly, bit for bit. + +Only --num-rollout is cut down, 10000 -> 2: one weight-sync round trip is +enough to catch drift in the post-update rollout and the second optimizer step. +--cuda-visible-devices "" unpins the recipe's default 0,1,2,3 so it inherits +the runner's GPU set (the 5gpu runner exposes GPUs 3-7). + +What this test uniquely guards is the bitwise train<->rollout parity the +cosmos3_bitwise rollout patch group + --sglang-lora-merge-mode dynamic +establish: train/model_output_mean_abs_diff / train/model_output_rel_max +compare the raw DiT outputs between engine and trainer — with parity both are +exactly 0. The recipe recomputes old log-probs at ingestion, so the log_prob +series guard the trainer-side pipeline rather than cross-side parity. +""" + +from tests.ci.e2e_metrics_registry import register_e2e_ci + +register_e2e_ci( + est_time=4800, + suite="stage-c-5-gpu-h200", + script="scripts/run_diffusion_grpo_cosmos3_pickscore_t2i_4gpu_bitwise.py", + args=["--num-rollout", "2", "--cuda-visible-devices", ""], + labels=["e2e"], + metrics=[ + "rollout/reward/raw_num_samples", + "rollout/reward/raw_mean", + "rollout/reward/raw_median", + "rollout/reward/raw_std", + "train/log_prob_old_idx_0", + "train/log_prob_new_idx_0", + "train/log_prob_mean_abs_diff", + "train/model_output_mean_abs_diff", + "train/model_output_rel_max", + "train/grad_norm", + ], +) From 4937d87516bfdd2326aafe9e5eaab95057959a7e Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Tue, 18 Aug 2026 14:24:25 +0000 Subject: [PATCH 4/5] ci: record the cosmos3 t2i e2e standard Recorded offline on 4xH200 in a replica of the CI image environment (torch 2.11.0+cu129, sglang main 0065fbfae1, diffusers 0.39.0): the model_output parity series are exactly 0 on both optimizer steps. Co-Authored-By: Claude Fable 5 --- ...test_cosmos3_pickscore_grpo_t2i_4xGPU.json | 108 ++++++++++++++++++ 1 file changed, 108 insertions(+) create mode 100644 tests/ci/fixtures/e2e_standards/test_cosmos3_pickscore_grpo_t2i_4xGPU.json diff --git a/tests/ci/fixtures/e2e_standards/test_cosmos3_pickscore_grpo_t2i_4xGPU.json b/tests/ci/fixtures/e2e_standards/test_cosmos3_pickscore_grpo_t2i_4xGPU.json new file mode 100644 index 00000000..0c80581e --- /dev/null +++ b/tests/ci/fixtures/e2e_standards/test_cosmos3_pickscore_grpo_t2i_4xGPU.json @@ -0,0 +1,108 @@ +{ + "meta": { + "commit": "local", + "source": "test_cosmos3_pickscore_grpo_t2i_4xGPU.py" + }, + "metrics": { + "rollout/reward/raw_mean": [ + [ + 0, + 0.7647189497947693 + ], + [ + 1, + 0.7742270827293396 + ] + ], + "rollout/reward/raw_median": [ + [ + 0, + 0.7654154300689697 + ], + [ + 1, + 0.7737646102905273 + ] + ], + "rollout/reward/raw_num_samples": [ + [ + 0, + 768.0 + ], + [ + 1, + 768.0 + ] + ], + "rollout/reward/raw_std": [ + [ + 0, + 0.06847453862428665 + ], + [ + 1, + 0.061292387545108795 + ] + ], + "train/grad_norm": [ + [ + 1.0, + 5.782852167612873e-05 + ], + [ + 2.0, + 3.818701588897966e-05 + ] + ], + "train/log_prob_mean_abs_diff": [ + [ + 1.0, + 0.0 + ], + [ + 2.0, + 0.0 + ] + ], + "train/log_prob_new_idx_0": [ + [ + 1.0, + -0.2110141608864069 + ], + [ + 2.0, + -0.3804309892778595 + ] + ], + "train/log_prob_old_idx_0": [ + [ + 1.0, + -0.2110141608864069 + ], + [ + 2.0, + -0.3804309892778595 + ] + ], + "train/model_output_mean_abs_diff": [ + [ + 1.0, + 0.0 + ], + [ + 2.0, + 0.0 + ] + ], + "train/model_output_rel_max": [ + [ + 1.0, + 0.0 + ], + [ + 2.0, + 0.0 + ] + ] + } +} From 2976fe15e8cce64212c99cdb211881fe7c55775c Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Tue, 18 Aug 2026 16:30:46 +0000 Subject: [PATCH 5/5] fix(cosmos3): request the tuple form from the diffusers f53d5520 forward Cosmos3OmniTransformer.forward defaults to a Cosmos3OmniTransformerOutput since f53d5520 (the requirements pin); unpacking that object yields only the non-None fields and T2I has neither sound nor action, so the 3-tuple unpack got 1 value. return_dict=False returns the unchanged tuple. Co-Authored-By: Claude Fable 5 --- miles/backends/fsdp_utils/configs/cosmos3.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/miles/backends/fsdp_utils/configs/cosmos3.py b/miles/backends/fsdp_utils/configs/cosmos3.py index 97d6f304..12d0a2ff 100644 --- a/miles/backends/fsdp_utils/configs/cosmos3.py +++ b/miles/backends/fsdp_utils/configs/cosmos3.py @@ -154,6 +154,9 @@ def _packed_forward( vision_mse_loss_indexes=vision_sequence_indexes, vision_timesteps=torch.full((num_vision_tokens,), timestep, device=device, dtype=torch.float32), vision_noisy_frame_indexes=[torch.arange(latent_t, dtype=torch.long, device=device)], + # diffusers f53d5520 defaults to a Cosmos3OmniTransformerOutput; the + # tuple form is unchanged behind return_dict=False. + return_dict=False, ) return preds_vision[0]