[model,test,docs] feat: add FLUX.3 adapter

[Jayce-Ping]

Summary

Add Flux3Adapter (model_type: "flux3") supporting Black Forest Labs' newly released FLUX.3 Image Klein family (black-forest-labs/flux-3-image-klein-base and black-forest-labs/flux-3-image-klein) from huggingface/diffusers@7ecb1e1b4fe500be7fb7381a7bb79a24d7fa3941 (huggingface/diffusers#15046).

Key Changes

• Submodule Bump (diffusers): Update the bundled diffusers submodule to 7ecb1e1b4fe500be7fb7381a7bb79a24d7fa3941 (Flux3Transformer2DModel, Flux3ModularPipeline, Flux3DistilledModularPipeline, Flux3ImageProcessor). The new pin requires Python 3.11+ and huggingface-hub>=1.32; the README install note and .agents/knowledge/dependencies.md say so, and QWEN_IMAGE_21_DIFFUSERS_COMMIT now names the same pin (the Qwen-Image 2.1 functions it documents as copied are unchanged between the two commits).
• Component Runtime (Flux3Adapter):
  • Wrap Flux3ModularPipeline / Flux3DistilledModularPipeline with ModularPipelineRuntime, rewrite local snapshot ComponentSpec paths when loading from disk, and declare component_load_dtype_defaults = {"transformers": torch.bfloat16, "text_encoders": torch.bfloat16, "vae": torch.float32}.
  • Opt into supports_fsdp2_cpu_efficient_loading = True and ddp_find_unused_parameters = True (required because mode_blocks.image_cond receives no gradient during text-only T2I steps).
  • Override _gradient_checkpointing_units() to return txt_mode_blocks -> mode_blocks.image -> mode_blocks.image_cond -> single_transformer_blocks in exact forward execution order.
• Prompt & Condition Image Encoding:
  • Implement Qwen3-VL text encoding with the official 80-token window right-padding recipe, explicit padding position_ids = 1, and 8-layer hidden-state stacking (4, 8, 12, 16, 20, 24, 28, 32) -> 20480 channels. Every text-encoding entry point (encode_prompt, inference, _inference, and the Qwen3-VL helpers) defaults to the official max_sequence_length=8192, so neither preprocessing nor raw-prompt rollouts truncate long captions.
  • Preserve per-sample 80-token windows without cross-sample padding pollution when batch prompts differ in window length.
  • Encode multi-reference condition images with AutoencoderKLFlux2 (sample_mode="argmax"), image_rope_scale = 0.5, and negative l coordinates -(1 + idx) (-1, -2, ...).
• Unified Base (True CFG) & Step-Distilled (Flow-Map) Forward / Inference:
  • Automatically detect flow_map=True (is_distilled) to pass target_timestep = t_next / 1000 and narrow _resolve_pipeline_io_contract() to NegativePromptPolicy.UNSUPPORTED.
  • Normalize img_ids, txt_ids, and image_cond_ids to 2D (seq_len, 4) for Flux3PosEmbed.
  • Run a batch whose prompts share one 80-token window count batched, and any other batch or any image-conditioned batch one sample at a time. Because a sample's bf16 result depends on its batch, the adapter declares requires_preserved_replay_batch_composition, and the optimize loop replays whole rollout packs (next item).
• Optimize-loop ordering (BaseTrainer): train.shuffle_samples now shuffles whole rollout packs, keeping each pack's samples and order, for adapters that declare requires_preserved_replay_batch_composition (Bagel, FLUX.3), so every micro-batch replays one rollout pack. Batch-invariant adapters still shuffle individual samples with the same permutation, and shuffle_samples: false still keeps the rollout order. The generation loop records the rollout packs of these adapters, and the ordering fails fast when a pack does not match per_device_batch_size or the samples changed after rollout; reward overlap shares that validation. Bagel's warning asking for shuffle_samples: false is removed, and the help text, guidance/workflow.md, guidance/algorithms.md, and the train-inference consistency notes describe the semantics.
• Offline SFT / DPO Support:
  • Extend src/flow_factory/models/flux/_output.py with encode_flux3_output_images (sample_mode="sample"), prepare_flux3_condition_latents, and standalone 2x2 patchify/unpatchify helpers.
  • Set offline_training_forward_overrides = MappingProxyType({"guidance_scale": 1.0}).
  • Register flux3 in the offline smoke catalog (text_to_image, image_to_image, multi_image_to_image).
• Examples & Docs:
  • Add examples/grpo/lora/flux3/default.yaml and update README.md, examples/README.md, guidance/datasets.md, and .agents/knowledge/architecture.md.

Fixes from GPU validation

Commit	Fix
[model] fix: default FLUX.3 encode_prompt to the official 8192-token limit	Preprocessing used the 512-token default; a 999-token caption got a 512-token window instead of the official 1040-token one.
[dataset,test] fix: register flux3 in the offline smoke catalog and README table check	tests/dataset/test_offline_smoke_profiles.py and tests/docs/test_sensenova_docs.py failed once flux3 was registered (both pass on main).
[model,docs,agents] fix: align Diffusers submodule references with the 7ecb1e1b pin	Stale Qwen-Image 2.1 pin and the undocumented Python 3.11 floor of the new submodule.

After validation, [model,test] fix: default every FLUX.3 entry point to the official 8192 tokens extends the 8192 default to inference, _inference, and the Qwen3-VL helpers. They still defaulted to 512 and passed it to encode_prompt, so raw-prompt rollouts (for example with enable_preprocess: false) truncated long captions.

[model] docs: describe the FLUX.3 tokenizer without reward decoding drops the tokenizer docstring's reference to reward-side prompt_ids decoding, which #249 removes; it matches #249's wording for Qwen-Image 2.1.

Fix from training validation: exact on-policy replay at per_device_batch_size > 1

A 32-GPU GRPO run of the GenEval example (per_device_batch_size: 8, from #248) logged on-policy ratios that rounded to 1.0000 but had ratio_std up to 1e-4, while Qwen-Image 2.1 under the same trainer logged exactly 1. inference() and forward() each ran a batch batched when all its prompts shared one 80-token window count and per sample otherwise. Group-contiguous rollout packs repeat one prompt and ran at batch size 8, while the shuffled optimize micro-batches mix window counts and replayed the same samples at batch size 1; bf16 GEMMs of different shapes round differently. The batched path only runs at per_device_batch_size > 1, which no earlier gate exercised: every GRPO gate config inherited per_device_batch_size: 1 from the shipped example.

Commit	Change
[model,test,agents] fix: run every FLUX.3 sample at transformer batch size 1 for exact replay	First fix: both phases ran every sample alone. Exact, but about 13% slower in rollout and 15% in optimization at per_device_batch_size: 8. Superseded by the next commit.
[trainer,model,test,docs,agents] feat: shuffle whole rollout packs for pack-dependent adapters	Final fix: FLUX.3 runs uniform-window packs batched again and declares requires_preserved_replay_batch_composition; the optimize loop shuffles the rollout packs of such adapters as whole units, so each micro-batch reproduces its rollout pack.

Verification

CPU

• Full suite before the final 8192 commit: pytest tests -> 3000 passed, 50 skipped, 0 failed with the submodule installed editable (pip install -e ./diffusers).
• Full suite on the final head (Linux CPU): pytest tests -> 3009 passed, 50 skipped, 0 failed. Merged with [prompt_enhance,model,cli,docs] feat: move prompt enhancement into ff-prompt-enhance #248 and [reward,trainer] fix: gather prompt strings for groupwise rewards #249, the full suite gives 3042 passed, 50 skipped, 0 failed.
• tests/trainers/test_optimize_sample_order.py checks the ordering against the exact torch.randperm oracle: batch-invariant adapters keep the previous per-sample permutation, pack-dependent adapters get a permutation of whole rollout packs (the case that fails on the previous trainer), single-sample packs reduce to the per-sample order, shuffle_samples: false keeps the rollout order, missing or mismatched packs and reordered samples fail fast, and the generation loop records packs only for pack-dependent adapters. test_rollout_packs_replay_exactly_through_their_rollout_path requires bit-identical replay of a uniform pack (batch size 2 in both phases) and of a mixed pack (batch size 1).
• scripts/validate_agent_harness.py, scripts/audit_docstrings.py (0 errors), and Black/isort on every changed Python file pass.

GPU (8x H20 97,871 MiB, driver 535.247.01)

Environment: Python 3.12, torch 2.10.0+cu129, transformers 5.14.1, accelerate 1.14.0, peft 0.20.0, deepspeed 0.18.3, huggingface-hub 1.32.0, diffusers 0.41.0.dev0 = submodule 7ecb1e1b. Weights: flux-3-image-klein-base@3f7fa5ee and flux-3-image-klein@33733762.

Tested commit: 516f9653. Like the tested commit, the final head defaults all five text-encoding entry points to 8192, now through a shared DEFAULT_MAX_SEQUENCE_LENGTH, and its test covers both the encode_prompt and raw-prompt inference paths. The parity script passes max_sequence_length explicitly, so its results do not depend on the default.

1. Real-weight 4-layer parity (1 GPU per variant; T2I and 2-reference I2I at 1024², base 28 steps CFG 3.5, distilled 16 steps)

• Config + components (L1/L2), both variants: scheduler (shift=6.93, no dynamic shifting), VAE scale, image-processor, and transformer configs identical; every transformer / Qwen3-VL / VAE weight bit-equal, including the 86 fp32 norm_q/norm_k tensors; prompt embeddings (80- and 160-token windows), a 999-token prompt at the new default (1040-token window), and the 8016 reference-image tokens are all bit-exact.
• End-to-end (L3/L4), strict mode, both variants: timesteps, initial noise, final latents, and decoded images are bit-exact (max abs diff 0) for T2I and I2I, without autocast and under bf16 autocast. Strict mode keeps trajectory latents in the transformer dtype (latent_storage_dtype: null) and makes the reference scheduler step with Flow-Factory's fp32 Euler arithmetic; everything else runs the official modular blocks.
  • In the autocast I2I case the reference images are encoded outside autocast, as dataset preprocessing does. Encoding inside autocast differs only because the shared FLUX.2/3 VAE helper normalizes in the VAE output dtype (as the official FLUX.2 pipelines do), while the official FLUX.3 step uses vae.dtype. The trainer never encodes conditions under autocast.
• Production defaults (informational): with fp16 trajectory storage and the shared scheduler's fp32 Euler (Diffusers computes dt * v in bf16), the first step already differs by one bf16 ulp and the trajectories drift: PSNR 37.1 / 43.6 dB (base T2I / +autocast), 53.4 / 43.8 dB (distilled T2I), 20.0 / 27.3 dB (base I2I), 30.1 / 29.9 dB (distilled I2I). Images keep the same subject and composition; the 28-step CFG trajectory amplifies the ulp differences. These settings are framework-wide, not FLUX.3-specific.

2. On-policy rollout/replay (real GRPO trainer)

A per-rank probe on the GRPO trainer compares every replayed log-prob with the stored rollout log-prob bit for bit, records the batch size of every FLUX.3 _forward call, and checks that every optimize micro-batch is one whole rollout pack. Each run keeps the default shuffle_samples: true and takes one optimizer step per epoch, so every replay is on-policy. Each ran 2 epochs on 8 GPUs with ZeRO-2 unless noted, per_device_batch_size: 8, 512², 10 steps, CPS, a PickScore reward, and 16 prompts x 8 per epoch; the second epoch rolls out after a real update. The final code was validated at 0b6b9b17, which differs from cf12271c by one CPU test.

Run	Before (ebac917a)	Per-sample (d7994e1f)	Final (cf12271c)
T2I, FLUX.3-upsampled GenEval prompts (several windows each)	421/768 log-probs differ, max 7.0e-6. Rollout ran at batch size 8; replay at 1 (432 calls) and 8 (42)	0/768	0/768. Rollout and replay at batch size 8; every micro-batch is a whole rollout pack, in shuffled pack order
Same, shuffle_samples: false	0/768: replay keeps the rollout packs
T2I, PickScore captions	140/768 differ, max 7.5e-6: some captions span two windows once the chat template is added	0/768	0/768, batch size 8
T2I, true CFG 3.5, gradient checkpointing		0/768	0/768, batch size 8
T2I, distilled flux-3-image-klein		0/768	0/768, batch size 8
I2I (ShareGPT-4o-Image-Mini, one reference), true CFG 3.5, gradient checkpointing		0/768	0/768; image-conditioned batches always run per sample
T2I, DDP		0/768	0/768, batch size 8
T2I, FSDP2		0/768	0/768, batch size 8
T2I, Flow-SDE (noise 0.7, one SDE step from [1, 2, 3])		0/256	0/256, batch size 8
The GenEval example from #248 as configured (group 18, 64 prompts, 2 optimizer steps per epoch, GenEval reward), 1 epoch	ratio_std up to 1e-4 in the 32-GPU run	0/1728 on-policy replays differ	0/1728; about 85% of sample forwards batched in both phases (packs that straddle two prompt groups run per sample). The trainer logs ratio_min=1 ratio_max=1 ratio_std=0 on the first step; the second step replays pre-update rollouts by design.

• Exact-state resume with the final code (save_model_only: false, resume_type: state from checkpoint-1, the GenEval-prompt run above): under DDP and ZeRO-2 the first resumed epoch reproduces the source run exactly: the same pack order on every rank, bit-equal rollout and replay log-probs, and all 40 logged train metrics equal.
• Throughput (GenEval-prompt run, second epoch): rollout 27.7 s and optimization 15.5 s with pack shuffling, the same as batched replay with shuffle_samples: false (27.8 s, 15.8 s); the per-sample commit took 31.5 s and 18.1 s.

Before this fix, the 21 GRPO gate runs (T2I/I2I, base/distilled, every backend, all at per_device_batch_size: 1, including the 384 first-step replays of the shipped example) already gave max |exp(new_log_prob − old_log_prob) − 1| = 0. The example's second gradient step per epoch reuses pre-update rollouts by design and stays within clip_range (max |ratio − 1| = 1.3e-5, clip fraction 0). norm_q/norm_k stay fp32 under DDP and FSDP2 (see the ZeRO-2 note below). Under DDP, LoRA gradients reach all four block streams in I2I; in T2I mode_blocks.image_cond correctly receives none.

3. Backend x algorithm smoke matrix (2 GPUs each unless noted; all exit 0, losses finite)

Run	ZeRO-2	DDP	FSDP2
GRPO LoRA, T2I (PickScore)	pass	pass (+ eval)	pass, also with fsdp_cpu_ram_efficient_loading: true
GRPO LoRA, multi-reference I2I (ShareGPT-4o-Image-Mini)	pass	pass	pass (+ eval)
GRPO LoRA, distilled flux-3-image-klein, T2I and I2I		pass
SFT LoRA, 1 epoch	pass	pass	pass
Offline-DPO LoRA, 1 epoch	pass	pass	pass
examples/grpo/lora/flux3/default.yaml as shipped (8 GPUs, ZeRO-2, eval at 1024² on 128 prompts, 1 epoch)	pass (eval PickScore 0.7920, 2 optimizer steps, 1,367 s)

4. Checkpoint round-trip

• save_model_only: true: the checkpoint reloaded into a fresh real-weight adapter matches all 516/516 saved LoRA tensors bit-exactly, and a resume_type: lora GRPO run trains on-policy from it.
• save_model_only: false, resuming from checkpoint-1 with resume_type: state (DDP, ZeRO-2, FSDP2 for T2I; DDP for I2I): the first resumed epoch reproduces the source run exactly (all logged metrics and per-rank rollout/replay log-probs bit-equal), so weights, RNG, and sampler state are restored. Under DDP and FSDP2, the optimizer state before the first resumed step is bit-equal to the source run's; under ZeRO-2 (where the step runs inside accelerator.backward), the Adam step counter continues from the checkpoint. Later epochs differ by design: two fresh runs from the same seed already diverge after the first optimizer step (non-deterministic GPU backward), as a control run showed.

Known framework-level behavior (not introduced here)

• Under DeepSpeed bf16, DeepSpeedEngine casts the whole module with module.bfloat16(), so Diffusers fp32 islands (norm_q/norm_k here; Wan and MiniMax H3 have the same) become bf16 and torch.rms_norm falls back from the fused kernel with a dtype-mismatch warning. Rollout and replay stay consistent (ratio exactly 1). Restoring fp32 islands after DeepSpeed preparation belongs in the trainer and is left for a separate PR.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

An organization owner can choose the repository organization as the billing source for reviews requested by members with a paid Copilot license. Learn about Copilot code review billing.