feat(qwen-image): bitwise train<->rollout parity via autocast-parity norms and a pinned flash kernel

[zhihengy]

Summary

Full-stack precision check of the Qwen-Image GRPO pipeline (FSDP2 train actor vs sglang-d rollout engine) with the dumper toolchain, then fixes until every paired module output is bit-exact: 27 checkpoints × 8 paired forwards (2 samples × 2 SDE steps × cond/uncond, LoRA r64 at step 0) all 0-ULP, including the root DiT output; the full record sweep shows 106/106 same-dtype pairs identical.

Starting point was final-output rel err ≈ 4.5e-3 per forward. Three independent root causes, in execution order:

1. Fused-norm rounding points (first divergence, every *_norm1/*_norm2 site, rel ~1.4e-3). The sgld LayerNormScaleShift patch rounds the LN output to bf16 and modulates in bf16 eager — the pre-refactor(fsdp)(4.3/5): family-declared input dtype policy at the model boundary #103 train semantics. Since refactor(fsdp)(4.3/5): family-declared input dtype policy at the model boundary #103 the train side runs under ambient autocast: nn.LayerNorm computes and returns fp32, (1 + scale) rounds once in bf16, the modulation promotes to fp32, and the single bf16 rounding happens at the next matmul input. One extra rounding on the rollout side broke parity.
   Fix: new rollout patch group qwen_image (monkey_patches/patch_qwen_image_autocast_norms.py) re-patching LayerNormScaleShift + ScaleResidualLayerNormScaleShift with the autocast rounding points. It is applied after sgld (family group appends after the CLI group) and declared per-family via QwenImageTrainPipelineConfig.rollout_patch_group, so the recorded SD3/Wan parity standards (which bake in the old sgld semantics) are untouched.

2. Attention kernel mismatch. Offline arbitration on dumped q/k/v proved the rollout (--sglang-attention-backend torch_sdpa) executes SDPA flash bitwise, while the train-side diffusers default dispatch picks a different kernel (max |Δ| ~2.5e-1 vs flash on H200).
   Fix: pin --fsdp-attention-backend _native_flash in the aligned pickscore recipe (verified fwd+bwd works under --deterministic-mode).

3. Identity attention mask blocks flash. collate_cond_for_sample_batch always passed encoder_hidden_states_mask; with micro-batch 1 it is all-True (identity), but any non-None mask makes the flash path raise — and the rollout side always runs mask-free batch-1 requests.
   Fix: drop the mask when mask.all() (bitwise no-op, unlocks the flash path and matches rollout dispatch).

Pre-existing pieces this rests on: _rebuild_pos_embed_freqs_on_cuda, the sgld op patches, and #92's exact-sigma replay. No dependency on #106 (the Qwen parallel plan stays empty).

Verification

• Reduced GRPO smoke (1 prompt × 2 samples, 8 denoising steps, SDE window [3,5), true-CFG 4.0, 512×512, fp32 master / bf16 forward, micro-batch 1, CFG unbatched, grad-ckpt off, 1 GPU train + 1 GPU rollout): prep_align_qwenimage.py validate → 216/216 bit-exact across all 8 paired forwards; root output 0-ULP.
• Forwards pair on a bit-exact bf16 latent anchor; the only non-identical records left are autocast fp32-vs-bf16 boundary dtypes (equal after bf16 rounding) and the RoPE table representation (real cos/sin vs complex — same rotation, proven by bit-exact to_q/norm_q/attn outputs).
• tests/fast/backends/sglang_diffusion_utils/test_rollout_patch_groups.py passes (extended to assert the new group registers); patch-order smoke: qwen_image applied after sgld wins on both norm classes.

Test plan

• pytest tests/fast/backends/sglang_diffusion_utils/test_rollout_patch_groups.py
• 2-GPU dump run of the reduced aligned recipe: all paired module outputs bit-exact (before: 4.5e-3 rel at root output)
• Full 5-GPU run-diffusion-grpo-pickscore-5gpu-flowgrpo-aligned.sh sanity run with the pinned flash backend

Made with Cursor