From 47fd1c1b8b6f3b8c3bc84c2dd36ed0f2af158246 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Tue, 11 Aug 2026 00:48:17 +0000 Subject: [PATCH 1/4] feat(qwen-image): bitwise train<->rollout parity via the qwen_image rollout patch group --- .../backends/fsdp_utils/configs/qwen_image.py | 12 +- .../monkey_patches/__init__.py | 26 +-- .../monkey_patches/_common.py | 7 - .../patch_layernorm_scale_shift.py | 24 --- .../monkey_patches/patch_mul_add.py | 22 --- .../monkey_patches/patch_qk_norm_rope.py | 59 ------ .../monkey_patches/patch_qwen_image.py | 177 ++++++++++++++++++ .../monkey_patches/patch_rmsnorm.py | 37 ---- .../patch_scale_residual_layernorm.py | 34 ---- miles/utils/arguments.py | 5 +- ...on_grpo_pickscore_5gpu_flowgrpo_aligned.py | 2 +- ...ion_grpo_qwen_image_bitwise_parity_4gpu.py | 128 +++++++++++++ .../test_rollout_patch_groups.py | 13 +- 13 files changed, 326 insertions(+), 220 deletions(-) delete mode 100644 miles/backends/sglang_diffusion_utils/monkey_patches/_common.py delete mode 100644 miles/backends/sglang_diffusion_utils/monkey_patches/patch_layernorm_scale_shift.py delete mode 100644 miles/backends/sglang_diffusion_utils/monkey_patches/patch_mul_add.py delete mode 100644 miles/backends/sglang_diffusion_utils/monkey_patches/patch_qk_norm_rope.py create mode 100644 miles/backends/sglang_diffusion_utils/monkey_patches/patch_qwen_image.py delete mode 100644 miles/backends/sglang_diffusion_utils/monkey_patches/patch_rmsnorm.py delete mode 100644 miles/backends/sglang_diffusion_utils/monkey_patches/patch_scale_residual_layernorm.py create mode 100644 scripts/run_diffusion_grpo_qwen_image_bitwise_parity_4gpu.py diff --git a/miles/backends/fsdp_utils/configs/qwen_image.py b/miles/backends/fsdp_utils/configs/qwen_image.py index 18268af5..466602c0 100644 --- a/miles/backends/fsdp_utils/configs/qwen_image.py +++ b/miles/backends/fsdp_utils/configs/qwen_image.py @@ -156,18 +156,16 @@ def collate_cond_for_sample_batch( 1 ) # (M, max_len) bool + # Collapse an all-true mask to None so training hits the same mask-less SDPA kernel as rollout. + mask_or_none = None if bool(mask.all()) else mask + return { "encoder_hidden_states": encoder_hidden_states, - "encoder_hidden_states_mask": mask, + "encoder_hidden_states_mask": mask_or_none, "txt_seq_lens": seq_lens, "img_shapes": img_shapes, } - def maybe_legacy_window_pad_len(self, conds: list) -> int | None: - # LEGACY 2D parity: reproduce the legacy whole-window cond pad width. TODO: remove with legacy 2D path. - lens = [int(c.txt_seq_lens[0]) for c in conds if c is not None and c.txt_seq_lens] - return max(lens) if lens else None - def cfg_combine( self, noise_pred_pos: torch.Tensor, @@ -186,7 +184,7 @@ def cfg_combine( combined = noise_pred_neg + scale * (noise_pred_pos - noise_pred_neg) if true_cfg_scale is not None and true_cfg_scale > 1.0: pos_norm = torch.norm(noise_pred_pos, dim=-1, keepdim=True) - combined_norm = torch.norm(combined, dim=-1, keepdim=True) + combined_norm = torch.norm(combined, dim=-1, keepdim=True).clamp_min(1e-12) combined = combined * (pos_norm / combined_norm) return combined diff --git a/miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py b/miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py index e377f75e..1d085a7f 100644 --- a/miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py +++ b/miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py @@ -5,11 +5,7 @@ the sglang scheduler grandchild (spawn: fresh imports) re-reads it and applies those groups before model construction. -- ``sgld``: diffusers / SD3 op parity (RMSNorm, LayerNormScaleShift, MulAdd, - ...). Op-layer patches: they apply to every sgl-d DiT built from these - generic classes. Attention is NOT patched: overriding USPAttention.forward - breaks bitwise SP-invariance (kernel choice depends on head/batch shape) — - align the attention kernel via the attention-backend selection instead. +- ``qwen_image``: bitwise train<->rollout parity for the Qwen-Image DiT. - ``ltx``: LTX rollout cond kwargs + AV cross-off (video-only train parity). Patch modules are imported inside ``apply_*`` only, so CPU-only Ray actors @@ -22,7 +18,7 @@ import os from collections.abc import Callable -# Comma-separated group names selected by the engine parent, e.g. "sgld,ltx". +# Comma-separated group names selected by the engine parent, e.g. "qwen_image". ROLLOUT_PATCH_GROUPS_ENV = "MILES_ROLLOUT_PATCH_GROUPS" _ROLLOUT_PATCH_APPLIERS: dict[str, Callable[[], None]] = {} @@ -38,21 +34,11 @@ def wrapper(fn: Callable[[], None]) -> Callable[[], None]: return wrapper -@register_rollout_patch_group("sgld") -def apply_sgld_monkey_patches() -> None: - from miles.backends.sglang_diffusion_utils.monkey_patches import ( - patch_layernorm_scale_shift, - patch_mul_add, - patch_qk_norm_rope, - patch_rmsnorm, - patch_scale_residual_layernorm, - ) +@register_rollout_patch_group("qwen_image") +def apply_qwen_image_rollout_patches() -> None: + from miles.backends.sglang_diffusion_utils.monkey_patches import patch_qwen_image - patch_rmsnorm.apply() - patch_layernorm_scale_shift.apply() - patch_scale_residual_layernorm.apply() - patch_mul_add.apply() - patch_qk_norm_rope.apply() + patch_qwen_image.apply() @register_rollout_patch_group("ltx") diff --git a/miles/backends/sglang_diffusion_utils/monkey_patches/_common.py b/miles/backends/sglang_diffusion_utils/monkey_patches/_common.py deleted file mode 100644 index 49f7a6ed..00000000 --- a/miles/backends/sglang_diffusion_utils/monkey_patches/_common.py +++ /dev/null @@ -1,7 +0,0 @@ -import torch - - -def ensure_broadcast(mod: torch.Tensor, ref: torch.Tensor) -> torch.Tensor: - if mod.dim() == ref.dim() - 1: - return mod.unsqueeze(-2) - return mod diff --git a/miles/backends/sglang_diffusion_utils/monkey_patches/patch_layernorm_scale_shift.py b/miles/backends/sglang_diffusion_utils/monkey_patches/patch_layernorm_scale_shift.py deleted file mode 100644 index e34ecfdc..00000000 --- a/miles/backends/sglang_diffusion_utils/monkey_patches/patch_layernorm_scale_shift.py +++ /dev/null @@ -1,24 +0,0 @@ -import torch - -from sglang.multimodal_gen.runtime.layers.layernorm import LayerNormScaleShift - -from miles.backends.sglang_diffusion_utils.monkey_patches._common import ensure_broadcast - - -def _patched_forward( - self, - x: torch.Tensor, - shift: torch.Tensor | None = None, - scale: torch.Tensor | None = None, -): - # diffusers sequence: LayerNorm(x) then (1+scale)*x + shift in bf16 eager. - normed = self.norm(x) - if shift is None and scale is None: - return normed - scale = ensure_broadcast(scale, normed) - shift = ensure_broadcast(shift, normed) - return normed * (1 + scale) + shift - - -def apply() -> None: - LayerNormScaleShift.forward = _patched_forward diff --git a/miles/backends/sglang_diffusion_utils/monkey_patches/patch_mul_add.py b/miles/backends/sglang_diffusion_utils/monkey_patches/patch_mul_add.py deleted file mode 100644 index 3ef8650e..00000000 --- a/miles/backends/sglang_diffusion_utils/monkey_patches/patch_mul_add.py +++ /dev/null @@ -1,22 +0,0 @@ -import torch - -from sglang.multimodal_gen.runtime.layers.elementwise import MulAdd - - -def _patched_forward( - self, - a: torch.Tensor, - b: torch.Tensor, - c: torch.Tensor, - k: int = 0, -): - # diffusers bf16 equivalent of the fused fp32 kernel: c + a*(k+b). - if b.dim() == 4: - num_frames = b.shape[1] - frame_seqlen = a.shape[1] // num_frames - return c + (a.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * (k + b)).flatten(1, 2) - return c + a * (k + b) - - -def apply() -> None: - MulAdd.forward = _patched_forward diff --git a/miles/backends/sglang_diffusion_utils/monkey_patches/patch_qk_norm_rope.py b/miles/backends/sglang_diffusion_utils/monkey_patches/patch_qk_norm_rope.py deleted file mode 100644 index 84c564c1..00000000 --- a/miles/backends/sglang_diffusion_utils/monkey_patches/patch_qk_norm_rope.py +++ /dev/null @@ -1,59 +0,0 @@ -import importlib - -import torch - -from sglang.multimodal_gen.runtime.layers import layernorm as _layernorm_mod - -# sgl-d DiT modules that import apply_qk_norm_with_optional_rope by name. -# Each one needs the name re-bound so monkey-patching layernorm alone isn't enough. -_REBIND_MODULES = ( - "sglang.multimodal_gen.runtime.models.dits.qwen_image", - "sglang.multimodal_gen.runtime.models.dits.flux", - "sglang.multimodal_gen.runtime.models.dits.flux_2", - "sglang.multimodal_gen.runtime.models.dits.zimage", -) - - -def _patched_apply_qk_norm_with_optional_rope( - q: torch.Tensor, - k: torch.Tensor, - q_norm, - k_norm, - head_dim: int, - cos_sin_cache=None, - *, - is_neox: bool = False, - positions=None, - position_offset: int = 0, - allow_inplace: bool = True, -): - # Replace sgl-d's fused qk-norm-rope CUDA kernel (which bypasses the - # patched RMSNorm.forward) with: patched q_norm/k_norm + diffusers' - # complex ROPE formula from apply_rotary_emb_qwen(use_real=False). - q_normed = q_norm(q) - k_normed = k_norm(k) - if cos_sin_cache is None: - return q_normed, k_normed - - # Layout: [cos_half | sin_half] along last dim; each half = head_dim/2. - half = cos_sin_cache.shape[-1] // 2 - freqs_cis = torch.complex(cos_sin_cache[..., :half], cos_sin_cache[..., half:]) - - def _apply(x: torch.Tensor) -> torch.Tensor: - x_c = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2)) - f = freqs_cis.unsqueeze(1).to(x.device) - if f.dim() < x_c.dim(): - f = f.unsqueeze(0) - return torch.view_as_real(x_c * f).flatten(3).type_as(x) - - return _apply(q_normed), _apply(k_normed) - - -def apply() -> None: - _layernorm_mod.apply_qk_norm_with_optional_rope = _patched_apply_qk_norm_with_optional_rope - for mod_path in _REBIND_MODULES: - try: - mod = importlib.import_module(mod_path) - except ImportError: - continue - mod.apply_qk_norm_with_optional_rope = _patched_apply_qk_norm_with_optional_rope diff --git a/miles/backends/sglang_diffusion_utils/monkey_patches/patch_qwen_image.py b/miles/backends/sglang_diffusion_utils/monkey_patches/patch_qwen_image.py new file mode 100644 index 00000000..ecd7e69d --- /dev/null +++ b/miles/backends/sglang_diffusion_utils/monkey_patches/patch_qwen_image.py @@ -0,0 +1,177 @@ +"""Qwen-Image rollout patches: make the sgl-d forward bitwise-equal to the diffusers/PEFT train forward.""" + +import torch +import torch.nn.functional as F +from torch.distributed.tensor import DTensor + +from sglang.multimodal_gen.runtime.layers import layernorm as layernorm_mod +from sglang.multimodal_gen.runtime.layers.elementwise import MulAdd +from sglang.multimodal_gen.runtime.layers.layernorm import ( + LayerNormScaleShift, + RMSNorm, + ScaleResidualLayerNormScaleShift, +) +from sglang.multimodal_gen.runtime.layers.lora import linear as lora_linear +from sglang.multimodal_gen.runtime.models.dits import qwen_image as qwen_image_mod + +_orig_split_seqs = qwen_image_mod.split_seqs +_orig_set_lora_weights = lora_linear.BaseLayerWithLoRA.set_lora_weights + + +def _rmsnorm_forward(self, x: torch.Tensor, residual: torch.Tensor | None = None): + # diffusers' RMSNorm rounds to weight dtype BEFORE the weight mul; sgl-d keeps fp32 through it. + if not x.is_contiguous(): + x = x.contiguous() + orig_dtype = x.dtype + x_fp32 = x.to(torch.float32) + if residual is not None: + x_fp32 = x_fp32 + residual.to(torch.float32) + residual = x_fp32.to(orig_dtype) + variance = x_fp32.pow(2).mean(dim=-1, keepdim=True) + x_fp32 = x_fp32 * torch.rsqrt(variance + self.variance_epsilon) + out = x_fp32.to(orig_dtype) + if self.weight is not None: + out = out * self.weight + if residual is None: + return out + return out, residual + + +def _ensure_broadcast(mod: torch.Tensor, ref: torch.Tensor) -> torch.Tensor: + if mod.dim() == ref.dim() - 1: + return mod.unsqueeze(-2) + return mod + + +def _fp32_layer_norm(norm: torch.nn.Module, x: torch.Tensor) -> torch.Tensor: + # nn.LayerNorm exactly as train-side autocast runs it: fp32 in, fp32 out. + weight = norm.weight.float() if norm.weight is not None else None + bias = norm.bias.float() if norm.bias is not None else None + return F.layer_norm(x.float(), norm.normalized_shape, weight, bias, norm.eps) + + +def _layernorm_scale_shift_forward( + self, + x: torch.Tensor, + shift: torch.Tensor | None = None, + scale: torch.Tensor | None = None, +): + normed = _fp32_layer_norm(self.norm, x) + if shift is None and scale is None: + return normed.to(x.dtype) + scale = _ensure_broadcast(scale, normed) + shift = _ensure_broadcast(shift, normed) + # (1 + scale) rounds in bf16, the modulation promotes to fp32 -- the train-side autocast semantics. + out = normed * (1 + scale) + shift + return out.to(x.dtype) + + +def _scale_residual_layernorm_scale_shift_forward( + self, + residual: torch.Tensor, + x: torch.Tensor, + gate: torch.Tensor, + shift: torch.Tensor, + scale: torch.Tensor, +): + residual_out = residual + x * gate + normed = _fp32_layer_norm(self.norm, residual_out) + scale = _ensure_broadcast(scale, normed) + shift = _ensure_broadcast(shift, normed) + out = normed * (1 + scale) + shift + return out.to(x.dtype), residual_out + + +def _mul_add_forward(self, a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, k: int = 0): + # diffusers bf16 equivalent of the fused fp32 kernel. + return c + a * (k + b) + + +def _qk_norm_rope( + q: torch.Tensor, + k: torch.Tensor, + q_norm, + k_norm, + head_dim: int, + cos_sin_cache=None, + *, + is_neox: bool = False, + positions=None, + position_offset: int = 0, + allow_inplace: bool = True, +): + # Replace the fused qk-norm-rope CUDA kernel with the patched norms + diffusers' complex RoPE. + q_normed = q_norm(q) + k_normed = k_norm(k) + if cos_sin_cache is None: + return q_normed, k_normed + + half = cos_sin_cache.shape[-1] // 2 + freqs_cis = torch.complex(cos_sin_cache[..., :half], cos_sin_cache[..., half:]) + + def _apply(x: torch.Tensor) -> torch.Tensor: + x_c = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2)) + f = freqs_cis.unsqueeze(1).to(x.device) + if f.dim() < x_c.dim(): + f = f.unsqueeze(0) + return torch.view_as_real(x_c * f).flatten(3).type_as(x) + + return _apply(q_normed), _apply(k_normed) + + +def _contiguous_split_seqs(joint, prefix_len, local_pad, dim=1): + # batch>1 split views are strided; contiguize so the out-proj GEMMs match diffusers' flattened GEMM. + prefix, body = _orig_split_seqs(joint, prefix_len, local_pad, dim=dim) + return prefix.contiguous(), body.contiguous() + + +def _lora_delta(self, x: torch.Tensor) -> torch.Tensor: + # PEFT-ordered LoRA path: (x @ A.T) @ B.T, then scale. + lora_A, lora_B = self.lora_A, self.lora_B + if isinstance(lora_B, DTensor): + lora_B = lora_B.to_local() + lora_A = lora_A.to_local() + x_lora = x.to(dtype=lora_A.dtype) + delta = x_lora @ self.slice_lora_a_weights(lora_A.to(device=x.device)).T + delta = delta @ self.slice_lora_b_weights(lora_B.to(device=x.device)).T + if self.lora_alpha != self.lora_rank: + delta = delta * (self.lora_alpha / self.lora_rank) + if self.strength != 1.0: + delta = delta * self.strength + return delta + + +def _lora_base_forward(self, x: torch.Tensor): + # base(x) first (bias included, as PEFT does), then the unmerged delta; bf16 add order matters. + out, output_bias = self.base_layer(x) + if not self.merged and not self.disable_lora: + out = out + _lora_delta(self, x).to(dtype=out.dtype) + return out, output_bias + + +def _lora_nn_linear_forward(self, x: torch.Tensor): + out = self.base_layer(x) + if not self.merged and not self.disable_lora: + out = out + _lora_delta(self, x).to(dtype=out.dtype) + return out + + +def _set_lora_weights_unmerged(self, *args, **kwargs): + # Merging W' = W + scaling*(B@A) in bf16 rounds differently from PEFT's unmerged path. + kwargs["merge_weights"] = False + return _orig_set_lora_weights(self, *args, **kwargs) + + +def apply() -> None: + RMSNorm.forward = _rmsnorm_forward + LayerNormScaleShift.forward = _layernorm_scale_shift_forward + ScaleResidualLayerNormScaleShift.forward = _scale_residual_layernorm_scale_shift_forward + MulAdd.forward = _mul_add_forward + layernorm_mod.apply_qk_norm_with_optional_rope = _qk_norm_rope + qwen_image_mod.apply_qk_norm_with_optional_rope = _qk_norm_rope + qwen_image_mod.split_seqs = _contiguous_split_seqs + lora_linear.BaseLayerWithLoRA.forward = _lora_base_forward + lora_linear.RowParallelLinearWithLoRA.forward = _lora_base_forward + lora_linear.ColumnParallelLinearWithLoRA.forward = _lora_base_forward + lora_linear.LinearWithLoRA.forward = _lora_nn_linear_forward + lora_linear.BaseLayerWithLoRA.set_lora_weights = _set_lora_weights_unmerged diff --git a/miles/backends/sglang_diffusion_utils/monkey_patches/patch_rmsnorm.py b/miles/backends/sglang_diffusion_utils/monkey_patches/patch_rmsnorm.py deleted file mode 100644 index 86a67f69..00000000 --- a/miles/backends/sglang_diffusion_utils/monkey_patches/patch_rmsnorm.py +++ /dev/null @@ -1,37 +0,0 @@ -import torch - -from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm - - -def _patched_forward( - self, - x: torch.Tensor, - residual: torch.Tensor | None = None, -): - # diffusers' RMSNorm rounds to weight dtype BEFORE the weight mul, so the - # mul runs bf16*bf16. sgl-d's default keeps fp32 through the weight mul. - if not x.is_contiguous(): - x = x.contiguous() - orig_dtype = x.dtype - - x_fp32 = x.to(torch.float32) - if residual is not None: - x_fp32 = x_fp32 + residual.to(torch.float32) - residual = x_fp32.to(orig_dtype) - - variance_size_override = getattr(self, "variance_size_override", None) - x_var = x_fp32 if variance_size_override is None else x_fp32[..., :variance_size_override] - variance = x_var.pow(2).mean(dim=-1, keepdim=True) - x_fp32 = x_fp32 * torch.rsqrt(variance + self.variance_epsilon) - - out = x_fp32.to(orig_dtype) - if self.weight is not None: - out = out * self.weight - - if residual is None: - return out - return out, residual - - -def apply() -> None: - RMSNorm.forward = _patched_forward diff --git a/miles/backends/sglang_diffusion_utils/monkey_patches/patch_scale_residual_layernorm.py b/miles/backends/sglang_diffusion_utils/monkey_patches/patch_scale_residual_layernorm.py deleted file mode 100644 index dcb45694..00000000 --- a/miles/backends/sglang_diffusion_utils/monkey_patches/patch_scale_residual_layernorm.py +++ /dev/null @@ -1,34 +0,0 @@ -import torch -from sglang.multimodal_gen.runtime.layers.layernorm import ScaleResidualLayerNormScaleShift - -from miles.backends.sglang_diffusion_utils.monkey_patches._common import ensure_broadcast - - -def _patched_forward( - self, - residual: torch.Tensor, - x: torch.Tensor, - gate, - shift: torch.Tensor, - scale: torch.Tensor, -): - # diffusers sequence: residual + gate*x (bf16), then LayerNorm, then - # (1+scale)*x + shift. - if isinstance(gate, int): - assert gate == 1 - residual_out = residual + x - elif gate.dim() == 4: - num_frames = gate.shape[1] - frame_seqlen = x.shape[1] // num_frames - residual_out = residual + (x.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * gate).flatten(1, 2) - else: - residual_out = residual + x * gate - - normed = self.norm(residual_out) - scale = ensure_broadcast(scale, normed) - shift = ensure_broadcast(shift, normed) - return normed * (1 + scale) + shift, residual_out - - -def apply() -> None: - ScaleResidualLayerNormScaleShift.forward = _patched_forward diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 46563e48..8a98c534 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -367,9 +367,8 @@ def add_rollout_arguments(parser): default=None, help=( "Comma-separated rollout patch groups applied at sglang-d startup so its " - "forward is numerically aligned with the training side, e.g. 'sgld' " - "(diffusers op parity, small rollout perf hit) or 'ltx' " - "(see sglang_diffusion_utils/monkey_patches)." + "forward is numerically aligned with the training side, e.g. 'qwen_image' " + "or 'ltx' (see sglang_diffusion_utils/monkey_patches)." ), ) parser.add_argument( diff --git a/scripts/run_diffusion_grpo_pickscore_5gpu_flowgrpo_aligned.py b/scripts/run_diffusion_grpo_pickscore_5gpu_flowgrpo_aligned.py index a878e01f..850f1d24 100644 --- a/scripts/run_diffusion_grpo_pickscore_5gpu_flowgrpo_aligned.py +++ b/scripts/run_diffusion_grpo_pickscore_5gpu_flowgrpo_aligned.py @@ -65,7 +65,7 @@ def execute(args: ScriptArgs, data_dir: str) -> None: "--diffusion-step-strategy-path miles.rollout.step_strategy_hub.sde_window " "--diffusion-num-sde-steps 2 " "--diffusion-sde-window-range 3,5 " - "--rollout-patch-group sgld " + "--rollout-patch-group qwen_image " ) eval_args = ( diff --git a/scripts/run_diffusion_grpo_qwen_image_bitwise_parity_4gpu.py b/scripts/run_diffusion_grpo_qwen_image_bitwise_parity_4gpu.py new file mode 100644 index 00000000..066df603 --- /dev/null +++ b/scripts/run_diffusion_grpo_qwen_image_bitwise_parity_4gpu.py @@ -0,0 +1,128 @@ +"""Reproduce bitwise train<->rollout parity for Qwen-Image: model_output_*_abs_diff stays exactly 0. + +Every optimizer step is on-policy (num_steps_per_rollout=1) and the train micro-batch +reproduces one rollout microgroup exactly (micro_batch_size_sample == rollout_microgroup_size, +contiguous dp split), so the mask collapses to None and both sides run the same SDPA flash +kernel (--fsdp-attention-backend _native_flash, torch_sdpa rollout). + +Layout: 4 GPUs, train + rollout + pickscore reward all colocated. + +Usage: + python3 scripts/run_diffusion_grpo_qwen_image_bitwise_parity_4gpu.py +""" + +from dataclasses import dataclass + +import typer + +import miles.utils.external_utils.command_utils as U + +MODEL = "Qwen/Qwen-Image" +DATASET = "rockdu/miles-diffusion-datasets" +DATASET_SUBSET = "flowgrpo_pickscore" +WANDB_PROJECT = "miles-diffusion-grpo" + + +@dataclass +class ScriptArgs(U.ExecuteTrainConfig): + cuda_visible_devices: str = "4,5,6,7" + num_rollout: int = 3 + 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_qwen_image_bitwise_parity_4gpu_{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 8 " + "--n-samples-per-prompt 16 " + f"--num-rollout {args.num_rollout} " + "--num-steps-per-rollout 1 " + "--rollout-microgroup-size 8 " + "--diffusion-train-iter-order sample_major " + "--diffusion-num-steps 10 " + "--diffusion-guidance-scale 4.0 " + "--diffusion-true-cfg-scale 4.0 " + "--diffusion-noise-level 1.2 " + "--diffusion-height 512 " + "--diffusion-width 512 " + "--diffusion-step-strategy-path miles.rollout.step_strategy_hub.sde_window " + "--diffusion-num-sde-steps 2 " + "--diffusion-sde-window-range 3,5 " + "--rollout-patch-group qwen_image " + ) + + grpo_args = "--advantage-estimator grpo --globalize-reward-std --diffusion-clip-range 1e-4 " + + optimizer_args = "--lr 3e-4 --adam-beta2 0.999 --weight-decay 1e-4 " + + lora_args = "--use-lora --lora-ipc-weight-sync --lora-rank 64 --lora-alpha 128 --lora-init-weights gaussian " + + reward_args = ( + "--rm-type pickscore " + "--pickscore-num-workers 1 " + "--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 4 " + "--sglang-attention-backend torch_sdpa " + "--update-weight-buffer-size 2147483648 " + ) + + train_backend_args = ( + "--train-backend fsdp --fsdp-master-dtype fp32 --fsdp-reduce-dtype fp32 --diffusion-forward-dtype bf16 " + "--fsdp-attention-backend _native_flash " + ) + + perf_args = "--gradient-checkpointing --micro-batch-size-sample 8 --micro-batch-size-tstep 1 " + + misc_args = ( + "--actor-num-gpus-per-node 4 " + "--rollout-num-gpus 4 " + "--rollout-num-gpus-per-engine 1 " + "--num-gpus-per-node 4 " + "--colocate " + "--colocate-reward " + "--deterministic-mode " + "--diffusion-debug-mode " + ) + + U.execute_train( + train_args=( + f"{ckpt_args} {rollout_args} {grpo_args} {optimizer_args} " + f"{lora_args} {reward_args} {wandb_args} {sglang_args} {train_backend_args} {perf_args} " + f"{misc_args} {args.extra_args}" + ), + num_gpus_per_node=4, + config=args, + extra_env_vars={"PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True"}, + ) + + +@U.dataclass_cli +def main(args: ScriptArgs) -> None: + data_dir = prepare(args) + execute(args, data_dir) + + +if __name__ == "__main__": + typer.run(main) 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 b8fcd741..a45761f4 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 @@ -42,15 +42,16 @@ def test_unknown_group_fails_loud(self, monkeypatch): mp.apply_env_selected_rollout_patches() def test_builtin_group_registered(self): - # The decorator ran at import time for the in-repo group. - assert "sgld" in mp._ROLLOUT_PATCH_APPLIERS + # The decorator ran at import time for the in-repo groups. + assert "qwen_image" in mp._ROLLOUT_PATCH_APPLIERS + assert "ltx" in mp._ROLLOUT_PATCH_APPLIERS class TestValidateRolloutPatchGroups: # The arg-validation entry point behind --rollout-patch-group: - # --rollout-patch-group "sgld,ltx" ──► registered appliers ──► pass - # --rollout-patch-group "sgld,bogus" ─► "bogus" unregistered ─► ValueError + # --rollout-patch-group "qwen_image,ltx" ──► registered appliers ──► pass + # --rollout-patch-group "qwen_image,bogus" ─► "bogus" unregistered ─► ValueError def test_known_pass_unknown_raises(self): - mp.validate_rollout_patch_groups(["sgld", "ltx"]) + mp.validate_rollout_patch_groups(["qwen_image", "ltx"]) with pytest.raises(ValueError, match="Unknown rollout patch group"): - mp.validate_rollout_patch_groups(["sgld", "bogus"]) + mp.validate_rollout_patch_groups(["qwen_image", "bogus"]) From 8ff9f86013f75cd025ddd7d306a4bee27a3c5865 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Tue, 11 Aug 2026 02:56:47 +0000 Subject: [PATCH 2/4] restore legacy window pad and cond mask; rename parity script to max alignment Co-authored-by: Cursor --- miles/backends/fsdp_utils/configs/qwen_image.py | 10 ++++++---- ...n_diffusion_grpo_qwen_image_max_alignment_4gpu.py} | 11 +++++------ 2 files changed, 11 insertions(+), 10 deletions(-) rename scripts/{run_diffusion_grpo_qwen_image_bitwise_parity_4gpu.py => run_diffusion_grpo_qwen_image_max_alignment_4gpu.py} (89%) diff --git a/miles/backends/fsdp_utils/configs/qwen_image.py b/miles/backends/fsdp_utils/configs/qwen_image.py index 466602c0..d80f67cd 100644 --- a/miles/backends/fsdp_utils/configs/qwen_image.py +++ b/miles/backends/fsdp_utils/configs/qwen_image.py @@ -156,16 +156,18 @@ def collate_cond_for_sample_batch( 1 ) # (M, max_len) bool - # Collapse an all-true mask to None so training hits the same mask-less SDPA kernel as rollout. - mask_or_none = None if bool(mask.all()) else mask - return { "encoder_hidden_states": encoder_hidden_states, - "encoder_hidden_states_mask": mask_or_none, + "encoder_hidden_states_mask": mask, "txt_seq_lens": seq_lens, "img_shapes": img_shapes, } + def maybe_legacy_window_pad_len(self, conds: list) -> int | None: + # LEGACY 2D parity: reproduce the legacy whole-window cond pad width. TODO: remove with legacy 2D path. + lens = [int(c.txt_seq_lens[0]) for c in conds if c is not None and c.txt_seq_lens] + return max(lens) if lens else None + def cfg_combine( self, noise_pred_pos: torch.Tensor, diff --git a/scripts/run_diffusion_grpo_qwen_image_bitwise_parity_4gpu.py b/scripts/run_diffusion_grpo_qwen_image_max_alignment_4gpu.py similarity index 89% rename from scripts/run_diffusion_grpo_qwen_image_bitwise_parity_4gpu.py rename to scripts/run_diffusion_grpo_qwen_image_max_alignment_4gpu.py index 066df603..f5aa6f9d 100644 --- a/scripts/run_diffusion_grpo_qwen_image_bitwise_parity_4gpu.py +++ b/scripts/run_diffusion_grpo_qwen_image_max_alignment_4gpu.py @@ -1,14 +1,14 @@ -"""Reproduce bitwise train<->rollout parity for Qwen-Image: model_output_*_abs_diff stays exactly 0. +"""Maximally aligned Qwen-Image train<->rollout run: model_output_*_abs_diff in debug mode. Every optimizer step is on-policy (num_steps_per_rollout=1) and the train micro-batch reproduces one rollout microgroup exactly (micro_batch_size_sample == rollout_microgroup_size, -contiguous dp split), so the mask collapses to None and both sides run the same SDPA flash -kernel (--fsdp-attention-backend _native_flash, torch_sdpa rollout). +contiguous dp split). The remaining diff comes from the legacy whole-window cond pad and its +attention mask, which move the train forward off the rollout's mask-less SDPA flash kernel. Layout: 4 GPUs, train + rollout + pickscore reward all colocated. Usage: - python3 scripts/run_diffusion_grpo_qwen_image_bitwise_parity_4gpu.py + python3 scripts/run_diffusion_grpo_qwen_image_max_alignment_4gpu.py """ from dataclasses import dataclass @@ -37,7 +37,7 @@ def prepare(args: ScriptArgs) -> str: def execute(args: ScriptArgs, data_dir: str) -> None: - run_name = f"diffusion_grpo_qwen_image_bitwise_parity_4gpu_{U.create_run_id()}" + run_name = f"diffusion_grpo_qwen_image_max_alignment_4gpu_{U.create_run_id()}" ckpt_args = f"--hf-checkpoint {MODEL} --save {args.output_dir}/{run_name}/ckpt --save-interval 10 " @@ -90,7 +90,6 @@ def execute(args: ScriptArgs, data_dir: str) -> None: train_backend_args = ( "--train-backend fsdp --fsdp-master-dtype fp32 --fsdp-reduce-dtype fp32 --diffusion-forward-dtype bf16 " - "--fsdp-attention-backend _native_flash " ) perf_args = "--gradient-checkpointing --micro-batch-size-sample 8 --micro-batch-size-tstep 1 " From 93b0579514d50a6b1dcf60f97f40e16d80bac719 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Tue, 11 Aug 2026 02:59:37 +0000 Subject: [PATCH 3/4] fix import order in patch_qwen_image Co-authored-by: Cursor --- .../sglang_diffusion_utils/monkey_patches/patch_qwen_image.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/miles/backends/sglang_diffusion_utils/monkey_patches/patch_qwen_image.py b/miles/backends/sglang_diffusion_utils/monkey_patches/patch_qwen_image.py index ecd7e69d..856a010e 100644 --- a/miles/backends/sglang_diffusion_utils/monkey_patches/patch_qwen_image.py +++ b/miles/backends/sglang_diffusion_utils/monkey_patches/patch_qwen_image.py @@ -2,8 +2,6 @@ import torch import torch.nn.functional as F -from torch.distributed.tensor import DTensor - from sglang.multimodal_gen.runtime.layers import layernorm as layernorm_mod from sglang.multimodal_gen.runtime.layers.elementwise import MulAdd from sglang.multimodal_gen.runtime.layers.layernorm import ( @@ -13,6 +11,7 @@ ) from sglang.multimodal_gen.runtime.layers.lora import linear as lora_linear from sglang.multimodal_gen.runtime.models.dits import qwen_image as qwen_image_mod +from torch.distributed.tensor import DTensor _orig_split_seqs = qwen_image_mod.split_seqs _orig_set_lora_weights = lora_linear.BaseLayerWithLoRA.set_lora_weights From c021559bc46baba22c92ac21f0df0993d192d5dc Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Wed, 12 Aug 2026 01:21:06 +0000 Subject: [PATCH 4/4] gate the lora forward patches to tp_size==1 The PEFT-ordered lora forward adds the rank-local delta after base_layer() has already all-reduced (RowParallel) or all-gathered (ColumnParallel with gather_output=True), which silently drops the other ranks' delta shares under TP>1. Bitwise train<->rollout parity is a tp_size==1 property anyway, so dispatch back to sglang's native TP-aware lora forwards when the layer is actually sharded. Co-Authored-By: Claude Fable 5 --- .../monkey_patches/patch_qwen_image.py | 22 +++++++++++++++++-- 1 file changed, 20 insertions(+), 2 deletions(-) diff --git a/miles/backends/sglang_diffusion_utils/monkey_patches/patch_qwen_image.py b/miles/backends/sglang_diffusion_utils/monkey_patches/patch_qwen_image.py index 856a010e..547b7e5a 100644 --- a/miles/backends/sglang_diffusion_utils/monkey_patches/patch_qwen_image.py +++ b/miles/backends/sglang_diffusion_utils/monkey_patches/patch_qwen_image.py @@ -15,6 +15,8 @@ _orig_split_seqs = qwen_image_mod.split_seqs _orig_set_lora_weights = lora_linear.BaseLayerWithLoRA.set_lora_weights +_orig_column_parallel_lora_forward = lora_linear.ColumnParallelLinearWithLoRA.forward +_orig_row_parallel_lora_forward = lora_linear.RowParallelLinearWithLoRA.forward def _rmsnorm_forward(self, x: torch.Tensor, residual: torch.Tensor | None = None): @@ -155,6 +157,22 @@ def _lora_nn_linear_forward(self, x: torch.Tensor): return out +def _lora_column_parallel_forward(self, x: torch.Tensor): + # The PEFT-ordered path adds the rank-local delta after base_layer() has already + # all-gathered (gather_output=True), so it only holds at tp_size==1; bitwise parity + # is unattainable under TP anyway, so fall back to the native TP-aware forward. + if self.base_layer.tp_size > 1: + return _orig_column_parallel_lora_forward(self, x) + return _lora_base_forward(self, x) + + +def _lora_row_parallel_forward(self, x: torch.Tensor): + # Same constraint: base_layer() all-reduces before the rank-local delta is added. + if self.base_layer.tp_size > 1: + return _orig_row_parallel_lora_forward(self, x) + return _lora_base_forward(self, x) + + def _set_lora_weights_unmerged(self, *args, **kwargs): # Merging W' = W + scaling*(B@A) in bf16 rounds differently from PEFT's unmerged path. kwargs["merge_weights"] = False @@ -170,7 +188,7 @@ def apply() -> None: qwen_image_mod.apply_qk_norm_with_optional_rope = _qk_norm_rope qwen_image_mod.split_seqs = _contiguous_split_seqs lora_linear.BaseLayerWithLoRA.forward = _lora_base_forward - lora_linear.RowParallelLinearWithLoRA.forward = _lora_base_forward - lora_linear.ColumnParallelLinearWithLoRA.forward = _lora_base_forward + lora_linear.RowParallelLinearWithLoRA.forward = _lora_row_parallel_forward + lora_linear.ColumnParallelLinearWithLoRA.forward = _lora_column_parallel_forward lora_linear.LinearWithLoRA.forward = _lora_nn_linear_forward lora_linear.BaseLayerWithLoRA.set_lora_weights = _set_lora_weights_unmerged