diff --git a/miles/backends/fsdp_utils/configs/qwen_image.py b/miles/backends/fsdp_utils/configs/qwen_image.py index 37234935..73005559 100644 --- a/miles/backends/fsdp_utils/configs/qwen_image.py +++ b/miles/backends/fsdp_utils/configs/qwen_image.py @@ -192,7 +192,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 edbc3c97..d0ede949 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("wan") 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..c54aa732 --- /dev/null +++ b/miles/backends/sglang_diffusion_utils/monkey_patches/patch_qwen_image.py @@ -0,0 +1,186 @@ +"""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 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 +from torch.distributed.tensor import DTensor + +_orig_split_seqs = qwen_image_mod.split_seqs +_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): + # 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 _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 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_row_parallel_forward + lora_linear.ColumnParallelLinearWithLoRA.forward = _lora_column_parallel_forward + lora_linear.LinearWithLoRA.forward = _lora_nn_linear_forward 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/backends/sglang_diffusion_utils/monkey_patches/patch_wan_norm_ops.py b/miles/backends/sglang_diffusion_utils/monkey_patches/patch_wan_norm_ops.py index 17ef6a27..585f9a27 100644 --- a/miles/backends/sglang_diffusion_utils/monkey_patches/patch_wan_norm_ops.py +++ b/miles/backends/sglang_diffusion_utils/monkey_patches/patch_wan_norm_ops.py @@ -12,7 +12,11 @@ from sglang.multimodal_gen.runtime.layers.elementwise import MulAdd from sglang.multimodal_gen.runtime.layers.layernorm import LayerNormScaleShift, ScaleResidualLayerNormScaleShift -from miles.backends.sglang_diffusion_utils.monkey_patches._common import ensure_broadcast + +def _ensure_broadcast(mod: torch.Tensor, ref: torch.Tensor) -> torch.Tensor: + if mod.dim() == ref.dim() - 1: + return mod.unsqueeze(-2) + return mod def _layer_norm_f32(norm: torch.nn.LayerNorm, x_f32: torch.Tensor) -> torch.Tensor: @@ -26,8 +30,8 @@ def _lnss_forward(self, x: torch.Tensor, shift=None, scale=None): normed = _layer_norm_f32(self.norm, x.float()) if shift is None and scale is None: return normed.type_as(x) - scale = ensure_broadcast(scale, normed).float() - shift = ensure_broadcast(shift, normed).float() + scale = _ensure_broadcast(scale, normed).float() + shift = _ensure_broadcast(shift, normed).float() return (normed * (1 + scale) + shift).type_as(x) @@ -51,8 +55,8 @@ def _srlnss_forward(self, residual: torch.Tensor, x: torch.Tensor, gate, shift, normed = _layer_norm_f32(self.norm, residual_out.float()) if shift is None and scale is None: return normed.type_as(residual_out), residual_out - scale = ensure_broadcast(scale, normed).float() - shift = ensure_broadcast(shift, normed).float() + scale = _ensure_broadcast(scale, normed).float() + shift = _ensure_broadcast(shift, normed).float() return (normed * (1 + scale) + shift).type_as(residual_out), residual_out diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index deeb75f9..5451ad46 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( @@ -1509,6 +1508,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 "qwen_image" 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( + "qwen_image 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/scripts/run_diffusion_grpo_pickscore_5gpu_flowgrpo_aligned.py b/scripts/run_diffusion_grpo_pickscore_5gpu_flowgrpo_aligned.py index 07bd7298..7e357646 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 = ( @@ -79,7 +79,11 @@ def execute(args: ScriptArgs, data_dir: str) -> None: 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 " + lora_args = ( + "--use-lora --lora-ipc-weight-sync --lora-rank 64 --lora-alpha 128 --lora-init-weights gaussian " + # PEFT evaluates adapters unmerged; merging rounds differently in bf16. + "--sglang-lora-merge-mode dynamic " + ) reward_args = ( "--rm-type pickscore " @@ -114,6 +118,7 @@ def execute(args: ScriptArgs, data_dir: str) -> None: "--num-gpus-per-node 5 " "--colocate " "--deterministic-mode " + "--diffusion-debug-mode " ) U.execute_train( diff --git a/tests/ci/fixtures/e2e_standards/test_qwenimage_pickscore_grpo_5xGPU.json b/tests/ci/fixtures/e2e_standards/test_qwenimage_pickscore_grpo_5xGPU.json new file mode 100644 index 00000000..d9cce721 --- /dev/null +++ b/tests/ci/fixtures/e2e_standards/test_qwenimage_pickscore_grpo_5xGPU.json @@ -0,0 +1,156 @@ +{ + "meta": { + "commit": "473daaf3d656b75ec9facc7a34ee65ca6d19eed3", + "source": "test_qwenimage_pickscore_grpo_5xGPU.py" + }, + "metrics": { + "rollout/reward/raw_mean": [ + [ + 0, + 0.8414278030395508 + ], + [ + 1, + 0.8220380544662476 + ] + ], + "rollout/reward/raw_median": [ + [ + 0, + 0.839218258857727 + ], + [ + 1, + 0.8214308023452759 + ] + ], + "rollout/reward/raw_num_samples": [ + [ + 0, + 512.0 + ], + [ + 1, + 512.0 + ] + ], + "rollout/reward/raw_std": [ + [ + 0, + 0.06459800899028778 + ], + [ + 1, + 0.06872709095478058 + ] + ], + "train/grad_norm": [ + [ + 1.0, + 0.00023576825333293527 + ], + [ + 2.0, + 0.0002070947375614196 + ], + [ + 3.0, + 0.0001272122171940282 + ], + [ + 4.0, + 0.00021652203577104956 + ] + ], + "train/log_prob_mean_abs_diff": [ + [ + 1.0, + 3.059918526560068e-05 + ], + [ + 2.0, + 5.1719252951443195e-05 + ], + [ + 3.0, + 2.9540504328906536e-05 + ], + [ + 4.0, + 3.591692075133324e-05 + ] + ], + "train/log_prob_new_idx_0": [ + [ + 1.0, + -0.9191523548215628 + ], + [ + 2.0, + -0.9185248203575611 + ], + [ + 3.0, + -0.9183900794014335 + ], + [ + 4.0, + -0.9186643585562706 + ] + ], + "train/log_prob_old_idx_0": [ + [ + 1.0, + -0.9191298661753535 + ], + [ + 2.0, + -0.9184782188385725 + ], + [ + 3.0, + -0.9183700243011117 + ], + [ + 4.0, + -0.9186263224110007 + ] + ], + "train/model_output_mean_abs_diff": [ + [ + 1.0, + 0.01750823500333354 + ], + [ + 2.0, + 0.024276575219118968 + ], + [ + 3.0, + 0.017203001349116676 + ], + [ + 4.0, + 0.020700254826806486 + ] + ], + "train/model_output_rel_max": [ + [ + 1.0, + 0.2912037670612335 + ], + [ + 2.0, + 0.284492164850235 + ], + [ + 3.0, + 0.25374430418014526 + ], + [ + 4.0, + 0.16948756575584412 + ] + ] + } +} diff --git a/tests/e2e/short/test_qwenimage_pickscore_grpo_5xGPU.py b/tests/e2e/short/test_qwenimage_pickscore_grpo_5xGPU.py new file mode 100644 index 00000000..1b22a21f --- /dev/null +++ b/tests/e2e/short/test_qwenimage_pickscore_grpo_5xGPU.py @@ -0,0 +1,39 @@ +"""E2E: Qwen-Image PickScore GRPO (flow_grpo-aligned), 5-GPU (4 colocated FSDP +DP=4 + sglang rollout engines, 1 dedicated pickscore GPU) — runs the example +script'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, 400 -> 2: one weight-sync round trip is enough +to catch drift in the post-update rollout and the second optimizer step. + +The train/log_prob_* and train/model_output_* series record the +train<->rollout residual under the qwen_image patch group + +--sglang-lora-merge-mode dynamic (the recipe runs without +--diffusion-recompute-old-log-prob, so old log-probs come from the engine). +Exact zero is not yet expected for Qwen-Image; the strict comparison pins the +residual at its recorded magnitude, so a numerics change on either side — +engine kernels, patch group, LoRA path — moves these series and fails loudly. +""" + +from tests.ci.e2e_metrics_registry import register_e2e_ci + +register_e2e_ci( + est_time=1200, + suite="stage-c-5-gpu-h200", + script="scripts/run_diffusion_grpo_pickscore_5gpu_flowgrpo_aligned.py", + args=["--num-rollout", "2"], + 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", + ], +) 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..a7e7519f 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 @@ -43,15 +43,16 @@ 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 "qwen_image" in mp._ROLLOUT_PATCH_APPLIERS assert "wan" 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", "wan"]) + mp.validate_rollout_patch_groups(["qwen_image", "ltx", "wan"]) with pytest.raises(ValueError, match="Unknown rollout patch group"): - mp.validate_rollout_patch_groups(["sgld", "bogus"]) + mp.validate_rollout_patch_groups(["qwen_image", "bogus"])