feat(qwen-image): bitwise train<->rollout parity via the qwen_image rollout patch group

[zhihengy]

With the conditions below, model_output_max_abs_diff / model_output_mean_abs_diff between the training forward and the rollout snapshot are exactly 0.0 at every optimizer step, verified across 3 rollouts including LoRA IPC weight syncs. scripts/run_diffusion_grpo_qwen_image_max_alignment_4gpu.py reproduces this end to end on 4 GPUs (train + rollout + colocated pickscore reward, debug mode on).

Necessary conditions for bitwise equality

• --num-steps-per-rollout 1: only on-policy steps can match; later steps train on updated weights.
• The text attention mask must be None: collate collapses an all-true mask to None so both sides hit the same mask-less SDPA flash kernel; a non-trivial mask changes the kernel dispatch.
• No legacy whole-window cond pad: the maybe_legacy_window_pad_len override is removed; batch-local padding reproduces the rollout request widths.
• The train micro-batch must reproduce a rollout microgroup exactly: micro_batch_size_sample == rollout_microgroup_size, micro_batch_size_tstep == 1, and --train-dp-split-mode contiguous (the default; stride round-robins pairs across ranks, mixing prompt lengths into one micro-batch, which makes the mask non-trivial and changes cuBLAS kernel selection).
• --fsdp-attention-backend _native_flash: pins the train forward to the same SDPA flash kernel the rollout executes.
• Rollout tp_size == 1: under TP the sharded GEMMs and all-reduce summation order change bf16 rounding, so bitwise parity is a TP=1 property by construction. The patch group stays safe to enable at TP>1 (see the LoRA section), it just cannot be bitwise there.

Patches (--rollout-patch-group qwen_image, all in monkey_patches/patch_qwen_image.py)

• RMSNorm: round to weight dtype before the weight mul, as diffusers does.
• LayerNormScaleShift / ScaleResidualLayerNormScaleShift: fp32 LayerNorm, bf16 (1 + scale) rounding, fp32 modulation — the train-side autocast semantics.
• MulAdd: eager bf16 c + a * (k + b) instead of the fused fp32 kernel.
• qk-norm + RoPE: replace the fused CUDA kernel with the patched norms plus diffusers' complex-multiplication RoPE.
• split_seqs: contiguize the joint-attention output slices before to_out/to_add_out, so batch>1 out-projections run the same flattened GEMM as diffusers.
• LoRA: keep IPC-synced adapters unmerged and evaluate them with PEFT's exact arithmetic — see below.

LoRA: why patch instead of sglang's native unmerged mode, and the TP gate

sgl-d does support unmerged LoRA natively (--lora-merge-mode dynamic, with TP-aware unmerged forwards on every *WithLoRA layer). The patch group still overrides two things, for different reasons:

• set_lora_weights forced to merge_weights=False: the IPC weight-sync path miles uses (post_training/weights_updater.py) hardcodes merge_weights=True and never consults lora_merge_mode, so without the override every sync merges W + s·(B@A) in bf16, which rounds differently from PEFT's unmerged path. This override is a stopgap: the clean fix is upstream (make the sgl-d updater respect lora_merge_mode), after which it can be dropped — miles already forwards any --sglang-lora-merge-mode flag to ServerArgs.
• LoRA forward rewrites: sglang's native unmerged forwards are not bitwise-PEFT-equal even at TP=1 — BaseLayerWithLoRA/LinearWithLoRA are @torch.compiled (inductor fusion changes bf16 rounding) and RowParallelLinearWithLoRA computes (Wx + delta) + bias where PEFT computes (Wx + bias) + delta. The patched forwards run PEFT's exact ordering: base output with bias first, then the scaled delta.

The PEFT-ordered forward adds the rank-local delta after base_layer() has already all-reduced (RowParallel) or all-gathered (ColumnParallel with gather_output=True), which would silently drop the other ranks' delta shares under TP>1 (with the default targets, to_out.0 is RowParallel). The forward patches are therefore gated per layer: at base_layer.tp_size > 1 they dispatch back to sglang's native TP-aware unmerged forwards, which add the delta before the collective. TP>1 rollouts with the patch group + LoRA are numerically correct, just not bitwise-equal to the training forward (nothing can be, per the conditions above).

Other changes

• Collapse an all-true encoder_hidden_states_mask to None and drop the legacy whole-window pad override in the Qwen-Image train pipeline config.
• clamp_min(1e-12) on the CFG rescale norm, mirroring the engine.
• The sgld patch group is dropped (it only served Qwen-Image); --rollout-patch-group now offers qwen_image and ltx, one self-contained group per model. The aligned pickscore recipe switches to qwen_image accordingly and is otherwise unchanged.

[Rockdu]

LGTM, may refactor LoRA patch later

[Rockdu]

Let me add CI for qwen-image as well so we get everything e2e-verified

[Rockdu]

LGTM