1. Model introduction
FLUX 3 Image Klein is a flow-matching DiT for text-to-image generation and image editing: per-stream mode blocks for the image and text tokens, then joint single-stream blocks over[image | text], with 4-axis RoPE and per-token AdaLN
modulation. Miles supports RL on the text-to-image task only (--diffusion-task t2i); image-conditioned requests are
not supported. Rollouts run on sglang-diffusion’s native FLUX 3 Image Klein pipeline (Flux3ImagePipeline) at
1024×1024 in bf16; the trainer runs the MilesModelBackend package miles/backends/fsdp_utils/models/flux3_image/,
which builds only the text-to-image path.
2. Supported variants
The step-distilled Klein variant samples without CFG through a different denoise loop and is not supported.
The launchers’
--model-path option also takes a local package directory.
3. Family config
Frommiles/backends/fsdp_utils/configs/flux3_image.py:
The
flux3_image rollout patch group (miles/backends/sglang_diffusion_utils/monkey_patches/patch_flux3_image_cfg_off.py)
runs a conditional-only denoise loop when the CFG scale is 1 or less: one DiT forward per step, the rollout transition
hook intact, CFG-on requests unchanged.
4. Launch
Both recipes put training, rollout and the reward model on the same 8 GPUs (--colocate plus
--<reward>-reward-colocate, one reward worker per engine) and decode rollout responses in parallel parser actors
(--sglang-server-concurrency 4 --rollout-fetch-in-parser --rollout-parser-num-workers 32). The recipes target
Hopper; on SM10x Blackwell, append --extra-args "--fsdp-attention-backend _flash_4".
4.1 DiffusionNFT + PickScore, LoRA, 8 GPUs colocated
scripts/run_diffusion_nft_flux3_image_pickscore_8gpu.py — LoRA rank 32 / alpha 64 synced over CUDA IPC in
dynamic merge mode (the engine keeps the adapter unmerged, like the trainer), an EMA rollout policy and deterministic
ODE rollouts at 16 steps without CFG: NFT trains the unguided velocity, so the rollout samples it too. 8 prompts × 8
samples per rollout, lr 3e-4, gradient checkpointing.
Status: 📈 V — Verified
4.2 Flow-GRPO + HPSv2.1, full finetune, 8 GPUs colocated
scripts/run_diffusion_grpo_flux3_image_hps_8gpu.py — full finetune (no LoRA) with fp32 master and reduce, bf16
forward, lr 1e-5; full weights sync over CUDA IPC. 10 sampling steps with CFG 3.5 on both the rollout and the trainer,
Flow-SDE noise level 0.7 on one SDE step per rollout drawn from steps 1-2, clip 1e-4, no KL. Old log-probs come from
the rollout: with this noise level the train-rollout drift leaves the ratio at 1 (approx_kl at most ~1e-12, clipfrac 0).
16 prompts × 16 samples per rollout, gradient checkpointing.
Status: 📈 V — Verified
5. Reference results
Both runs use the default seed and the recipe defaults on one 8×H200 node, except that the Flow-GRPO run stops at 150 rollouts (--num-rollout 150; the recipe default is 200). Held-out numbers use the first 32 prompts
of each reward’s test split, one image per prompt at a fixed seed, at the training step count.
5.1 DiffusionNFT + PickScore
The PickScore reward (rollout/reward/raw_mean, PickScore / 26) rises from 0.769 (mean of the first 10 rollouts) to
0.857 (mean of the last 10) over 100 rollouts:
The trained model without CFG beats the base model without CFG on 32/32 prompts and the base model with CFG 3.5 on
27/32. It also drifts toward PickScore’s preferences: brighter, more saturated images, and on some prompts it adds or
beautifies people.
Case studies (base with CFG 3.5, base without CFG, NFT without CFG):
An impressionist portrait of a shar pei
5.2 Flow-GRPO + HPSv2.1, full finetune
The HPSv2.1 reward rises from 0.263 (mean of the first 10 rollouts) to 0.278 (mean of the last 10) over 150 rollouts; it is flat for the first ~20 rollouts, as expected of a full finetune at lr 1e-5 with clip 1e-4:
6. Pairs well with
- LoRA Training and Weight Sync — IPC adapter sync and merge modes.
- Rewards — PickScore and HPSv2.1.