diff --git a/miles/backends/fsdp_utils/configs/qwen_image.py b/miles/backends/fsdp_utils/configs/qwen_image.py index 18268af5..d80f67cd 100644 --- a/miles/backends/fsdp_utils/configs/qwen_image.py +++ b/miles/backends/fsdp_utils/configs/qwen_image.py @@ -186,7 +186,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..547b7e5a --- /dev/null +++ b/miles/backends/sglang_diffusion_utils/monkey_patches/patch_qwen_image.py @@ -0,0 +1,194 @@ +"""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_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): + # 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 _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_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 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_max_alignment_4gpu.py b/scripts/run_diffusion_grpo_qwen_image_max_alignment_4gpu.py new file mode 100644 index 00000000..f5aa6f9d --- /dev/null +++ b/scripts/run_diffusion_grpo_qwen_image_max_alignment_4gpu.py @@ -0,0 +1,127 @@ +"""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). 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_max_alignment_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_max_alignment_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 " + ) + + 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"])