[Rockdu]

What

• Pin diffusers==0.41.0 (PyPI release) in requirements.txt, replacing the git pin f53d5520 (the MiniMax-H3 merge commit on diffusers main, 2026-08-05).
• Train Krea-2 with diffusers 0.40's attention processor: SDPA takes the 12 K/V heads with enable_gqa instead of K/V repeated to all 48 heads.

Why

• Qwen-Image 2.1 ( Add Qwen-Image 2.1 family support and a 1-GPU PickScore Flow-GRPO recipe #256) trains QwenImage21Transformer2DModel from Add Qwen-Image 2.1 huggingface/diffusers#14804 (merged 2026-09-18), which the current pin predates, so the trainer cannot resolve the DiT class.
• 0.41.0 (2026-10-06) is the first release that contains both MiniMax H3 and Qwen-Image 2.1, so the "H3 ships in no release yet" git pin can go.
• The image already satisfies 0.41.0's requirements (huggingface-hub 1.33.0 >= 1.32.0, safetensors 0.8.0 >= 0.8.0); no other package moves with it.
• 0.41 changed Krea2AttnProcessor to repeat K/V before SDPA, so every layer saves 4x the K/V for backward. On the first CI run of the bump alone, the Krea-2 async proxy (no gradient checkpointing) kept every metric identical to the standard but its train steps went from ~75 s to 101-170 s, and its final dcp.save died with NCCL Error 1: unhandled cuda error; every other e2e kept its step time and passed bitwise. Between the two images only diffusers moved (same base image digest, no NCCL/torch/CUDA package change, same miles code path).

What changes for the existing families

Read from git diff f53d5520 v0.41.0 (198 commits), limited to the code miles trains through:

Family / module	Change	Expected e2e effect
attention_dispatch (_native_attention, _flash_attention, _wrapped_flash_attn_3)	byte-identical	none
SD3, Wan2.2, Qwen-Image	type hints, a TP plan, Neuron-only RoPE; CUDA math unchanged	none
Krea-2	GQA attention repeats K/V heads instead of passing enable_gqa=True to SDPA	metrics unchanged, 4x K/V saved for backward (see Why); reverted by the processor fix
MiniMax H3 DiT	dtype casts read get_parameter_dtype(module) instead of module.weight.dtype	possible drift
MiniMax H3 VAE (SFT encoder)	_keep_in_fp32_modules narrowed (decoder no longer pinned to fp32), explicit input casts	possible drift in H3 SFT
Cosmos3	decoder loop moved into _run_decoder_stack	none
LTX-2.3	trained through ltx-core, not diffusers	none

Every diffusers symbol miles imports (18, including the DiT classes resolved from each checkpoint's _class_name) imports on 0.41.0.

Validation

• CPU: import check above.

• GPU: this PR touches requirements.txt, so CI builds a pr-326 image; labelled run-ci-image to run every suite on it.

• Bump alone: stage-c-3 (SD3 NFT, SD3 OCR GRPO, H3 t2va GRPO, H3 t2va SFT) and stage-c-5 (Wan2.2, Cosmos3, LTX-2.3, Qwen-Image) match their standards bitwise; Krea-2 async failed as above.

• With the Krea-2 processor fix (run 39090578815): all 9 e2e match their standards bitwise, Krea-2 async included. Krea-2 async train step times:

  Image	Train step times	Result
  latest (diffusers f53d5520)	78.4 / 73.8 / 73.8 / 73.7 s	pass
  pr-326, bump only	100.6 / 116.4 / 170.3 / 119.6 s	final dcp.save NCCL error
  pr-326, bump + processor fix	74.6 / 73.6 / 73.7 / 73.7 s	pass

• On CPU the patched processor hands SDPA (q_heads, kv_heads) = (8, 2) with enable_gqa=True (0.41: (8, 8)), forward output bitwise equal.

• The rebuild also re-resolves unpinned transitive packages and the FA3 kernels.lock (latest was built 2026-09-27). Any drift is checked against .claude/skills/install-miles-diffusion/snapshot/ and A/B'd with only diffusers swapped before it is attributed to diffusers.

Follow-up

• Refresh .claude/skills/install-miles-diffusion/snapshot/ with refresh.sh from the rebuilt image after merge (the snapshot is a capture of the image, not hand-edited).

Files

• requirements.txt: diffusers git pin -> diffusers==0.41.0.
• miles/backends/fsdp_utils/configs/krea2.py: diffusers 0.40's Krea2AttnProcessor.__call__, installed after the model materializes (trainer only); drop it once diffusers hands SDPA grouped K/V again.

Checklist

• pre-commit run --all-files passes — run on the changed files only (passed), not on all files
• Added/updated tests for new behaviour — no tests added (dependency pin; covered by the existing suites)
• pytest -x is green — not run locally (no GPU/sglang env); full suites run in CI on the pr-<num> image
• If launch flags changed, python3 train.py --help still parses — N/A, no flags changed
• If a public flag was added, it appears in the CLI reference docs — N/A
• If an example was added, it has a real walkthrough — N/A