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

.agents/knowledge/architecture.md

@@ -130,6 +130,7 @@ extensions are `GRPOGuardTrainer → GRPOTrainer`, `DPPOTrainer → GRPOTrainer`
 | `flux1-kontext` | `Flux1KontextAdapter` | Image-to-Image |
 | `flux2` | `Flux2Adapter` | Text-to-Image & Image(s)-to-Image |
 | `flux2-klein` | `Flux2KleinAdapter` | Text-to-Image & Image(s)-to-Image |
+| `flux3` | `Flux3Adapter` | Text-to-Image & Image(s)-to-Image |
 | `qwen-image` | `QwenImageAdapter` | Text-to-Image |
 | `qwen-image-edit-plus` | `QwenImageEditPlusAdapter` | Image(s)-to-Image |
 | `qwen-image-2.1` | `QwenImage21Adapter` | Text-to-Image & Image(s)-to-Image |

.agents/knowledge/dependencies.md

@@ -80,6 +80,8 @@ The authoritative list is `pyproject.toml` `[project.dependencies]` (20+ package
 ### diffusers
 - Use the released `diffusers>=0.40.0` package as the authoritative API. The repository submodule
   may be used for upstream development, but must not silently override the declared runtime dependency.
+- `flux3` and `qwen-image-2.1` need the pinned submodule (`7ecb1e1b`), which requires Python 3.11+
+  and `huggingface-hub>=1.32`; released diffusers keeps Python 3.10 support for the other adapters.
 - Model adapters depend on specific pipeline classes from diffusers. Major version bumps may rename or remove pipeline classes.
 - `load_pipeline()` in each adapter returns a `DiffusionPipeline`-compatible object; breaking changes in diffusers' pipeline API require adapter updates.
 

.agents/knowledge/topics/train_inference_consistency.md

@@ -36,7 +36,7 @@ If rollout and training `forward()` diverge, `ratio` deviates from 1.0 at epoch
 4. **Model weight change**: EMA swap without restore between `inference()` and first `forward()` call.
 5. **Scheduler state mismatch**: `step_index` not matching (e.g., dual-scheduler models).
 6. **`num_inference_steps` changed**: invalidates sigma schedule, all trajectory timesteps are wrong.
-7. **Batch/pack composition mismatch (pack-dependent adapters)**: For adapters whose batched `forward()` is *pack-composition-dependent* (e.g. Bagel NaViT sequence packing, where a sample's linear-projection matmuls run over the concatenated `sum_seqlen` of the whole pack), bf16 rounding depends on *which* samples share the pack. If a training micro-batch packs a different sample set than the corresponding rollout pack, the on-policy `forward()` is no longer bit-identical -> `ratio != 1` (~1e-4) even though every stored argument matches. Per-sample (B=1) adapters are immune. The trigger is the optimize-time sample shuffle reordering `samples` before chunking into micro-batches.
+7. **Batch/pack composition mismatch (pack-dependent adapters)**: For adapters whose batched `forward()` is *pack-composition-dependent* (e.g. Bagel NaViT sequence packing, where a sample's linear-projection matmuls run over the concatenated `sum_seqlen` of the whole pack), bf16 rounding depends on *which* samples share the pack. If a training micro-batch packs a different sample set than the corresponding rollout pack, the on-policy `forward()` is no longer bit-identical -> `ratio != 1` (~1e-4) even though every stored argument matches. Per-sample (B=1) adapters are immune. The trigger is an optimize-time shuffle that reorders individual samples before chunking into micro-batches; adapters declare the dependence with `requires_preserved_replay_batch_composition`, and the optimize loop then shuffles whole rollout packs (see below).
 8. **Stochastic conditioning encoder recomputed per forward**: When `forward()` *rebuilds* its conditioning from raw inputs each call (rather than replaying a stored embedding) and that encoder is non-deterministic, rollout and training diverge. Bagel I2I is the canonical case: the condition-image VAE (`DiagonalGaussian`, default `sample=True` -> `mean + std*randn`) is encoded **once** at rollout but **re-encoded every training `forward()`** (`_forward_rebuild` -> `_build_gen_context`), so each draws different noise -> different KV context -> on-policy `ratio != 1` (~2e-4). T2I is immune (text-only context, deterministic). Fix: make the condition encode deterministic (`vae.reg.sample = False`, posterior mean) or cache the rollout encoding and replay it. (Only affects `vae.encode` of conditions; generation uses init noise and `vae.decode` is unaffected.) **Bagel applies this fix**: `pipeline.vae.reg.sample = False` in `BagelAdapter.load_pipeline()`.
 9. **Preference arms replayed with the wrong batch**: DPO shares forward-process noise across chosen/rejected states, but each arm still owns its conditioning batch. Both policy and reference forwards for the rejected state must receive `rejected_batch`; pairing it with `chosen_batch` evaluates the rejected trajectory under another conditioning context.
 10. **Condition realization redrawn per candidate**: Geometry-bound VAE conditions and stochastic
@@ -48,14 +48,18 @@ If rollout and training `forward()` diverge, `ratio` deviates from 1.0 at epoch
     boundary so mask/layout binding cannot drift between candidates.
 11. **Model-specific velocity direction assumed by a trainer**: Standard flow adapters predict noise-ward velocity (`noise - clean`), while MiniMax H3 predicts data-ward velocity (`clean - noise`). Any `x0` target must use `adapter.project_velocity_to_clean_state()` rather than spelling `xt - sigma * velocity` inside a trainer.
 12. **Different forward graph for the same inputs**: A prefix KV cache (prefill + cached decode) and a joint forward are mathematically equal but not bf16-equal. On Qwen-Image 2.1 they differ by 0.9–2.5% relative L2 in velocity (about 6e-6 in fp32), far outside `clip_range` and exact replay tolerances; grad mode alone does not change the result. Rollout and replay must build the same graph, so Qwen-Image 2.1 records `use_kv_cache` on the rollout sample and replay reads it back as a batch key.
+13. **Execution path chosen from the batch's contents**: An adapter that runs a uniform batch batched and a mixed batch one sample at a time computes the same sample with different GEMM shapes whenever its rollout pack and its training micro-batch differ in that respect, and bf16 rounding then differs. FLUX.3 did this with its 80-token prompt windows: group-contiguous rollout packs repeat one prompt and ran batched, while shuffled micro-batches mix window counts and ran per sample, so about half the replays differed by up to 7e-6 in log-prob. Prompts long enough to span two windows trigger it, including some short captions once the chat template is added. FLUX.3 therefore declares `requires_preserved_replay_batch_composition`, so replay keeps its rollout packs (item 7).
 
 ## Pack-composition-dependent adapters: `shuffle_samples`
 
-`samples` reaches `optimize()` in rollout-pack order (`generate_samples` appends each `inference` pack in order; `compute_advantages(store_to_samples=True)` attaches advantages in place without reordering), so `samples[i*B:(i+1)*B]` is exactly rollout pack `i`. Set `train.shuffle_samples: false` so the optimize loop keeps that order: each training micro-batch then packs the *same* samples as its rollout pack -> deterministic bf16 -> on-policy `ratio == 1`.
+`samples` reaches `optimize()` in rollout-pack order (`generate_samples` appends each `inference` pack in order; `compute_advantages(store_to_samples=True)` attaches advantages in place without reordering), so `samples[i*B:(i+1)*B]` is exactly rollout pack `i`. `train.shuffle_samples` keeps that pack structure whenever the adapter needs it:
 
-- Requires matched sampling and training `per_device_batch_size` (so contiguous chunks reproduce the rollout packs). These are the *same* field: rollout uses `training_args.per_device_batch_size` for the train DataLoader (`data_utils/loader.py`) and `optimize()` chunks by the same value, so they cannot drift unless a separate sampling batch size is introduced.
-- Default is `True` (per-inner-epoch shuffle); only pack-dependent adapters (Bagel) need `false`. The off-policy decorrelation cost is minor because the rollout order is already sampler-randomized. `BagelAdapter` warns at init if `shuffle_samples` is left `True`.
-- Wired via `BaseTrainer._order_samples_for_optimize(samples, inner_epoch)` (used by grpo/dppo/nft/awm/opd; dpo shuffles chosen/rejected pairs separately).
+- `true` (default) with a batch-invariant adapter: every inner epoch permutes individual samples.
+- `true` with an adapter that declares `requires_preserved_replay_batch_composition` (Bagel, FLUX.3): every inner epoch permutes whole rollout packs and keeps each pack's samples and order, so each training micro-batch replays exactly one rollout pack -> deterministic bf16 -> on-policy `ratio == 1`. For these adapters `generate_samples` always records the rollout packs in an `AcquisitionManifest`, and the ordering fails fast if a pack does not match `per_device_batch_size` or the samples changed after rollout (for example a custom `generate_samples()` that does not record boundaries).
+- `false`: the rollout order is kept for every adapter.
+- Rollout and optimize chunk by the same `training_args.per_device_batch_size` (the train DataLoader in `data_utils/loader.py` and `optimize()`), so a replayed pack has its rollout size unless a separate sampling batch size is introduced.
+- Shuffling packs instead of samples costs little decorrelation: the rollout order is already sampler-randomized, and with one optimizer step per epoch the accumulated gradient covers the same samples in any order.
+- Wired via `BaseTrainer._order_samples_for_optimize(samples, inner_epoch)` (used by grpo/dppo/nft/awm/opd; dpo shuffles chosen/rejected pairs separately because it compares policy and reference on the same micro-batch rather than against rollout log-probs). Reward/optimization overlap keeps requiring `shuffle_samples: false` and validates the same packs against its reward tiles.
 
 ## Where in Code
 
@@ -84,6 +88,14 @@ If rollout and training `forward()` diverge, `ratio` deviates from 1.0 at epoch
 - **Lesson**: Never flip the grad mode above code that decides from it whether a graph may outlive the call. A graph built under forced grad without a later backward is a leak whenever activation checkpointing saves non-tensor inputs that point back into it.
 - **Related Constraint**: N/A
 
+### Replay FLUX.3 rollout packs whole
+- **Date**: 2026-10-18
+- **Symptom**: FLUX.3 GRPO with `per_device_batch_size: 8` logged on-policy ratios that rounded to 1.0000 but were not exactly 1 (`ratio_std` up to 1e-4), while Qwen-Image 2.1 under the same trainer stayed exactly 1. An exact-replay probe found 421 of 768 replayed log-probs off by up to 7e-6 with GenEval prompts rewritten by the FLUX.3 upsampler, and 140 of 768 with PickScore captions.
+- **Root Cause**: `inference()` and `forward()` each chose batched execution for a batch whose prompts shared one 80-token window count and per-sample execution otherwise. Rollout packs of one repeated prompt ran at batch size 8; after the optimize-loop shuffle the same samples replayed at batch size 1, or the other way round. Keeping the rollout packs (`shuffle_samples: false`) made every replay exact, which isolated the pack composition.
+- **Fix**: FLUX.3 declares `requires_preserved_replay_batch_composition`, and `shuffle_samples` now shuffles a pack-dependent adapter's recorded rollout packs as whole units instead of individual samples, so every micro-batch replays one rollout pack in rollout order. FLUX.3 keeps batched execution for uniform-window packs. Running every sample alone was also exact but cost about 13% rollout and 15% optimization time at `per_device_batch_size: 8`.
+- **Lesson**: A GPU parity gate must include `per_device_batch_size > 1`, prompts of mixed lengths, and the default shuffle, and compare log-probs exactly rather than through the rounded ratio metric. Gates inherited from a `per_device_batch_size: 1` example never ran the batched path.
+- **Related Constraint**: Constraint #7
+
 ## Cross-refs
 
 - UP: [`constraints.md` #7](../constraints.md#7-coupled-vs-decoupled-paradigm), [Architecture Execution Pipelines](../architecture.md#execution-pipelines)

README.md

@@ -9,6 +9,17 @@
 
 # 🔥 News
 
+* **[2026-10-16]** **FLUX.3 Image Klein** support! Train both the base (`black-forest-labs/flux-3-image-klein-base`)
+  and step-distilled flow-map (`black-forest-labs/flux-3-image-klein`) text-to-image and multi-reference
+  image-to-image models with the [GRPO + LoRA recipe](examples/grpo/lora/flux3/default.yaml).
+  Until FLUX.3 reaches a released Diffusers package, install the pinned submodule
+  (it requires Python 3.11 or newer):
+```bash
+git submodule update --init
+pip install -e ./diffusers
+pip install -e .
+```
+
 * **[2026-09-20]** **Qwen-Image 2.1** support! Train the unified text-to-image and
   ordered multi-image editing model with the [GRPO + LoRA recipe](examples/grpo/lora/qwen_image_2_1/default.yaml).
   Until Qwen-Image 2.1 reaches a released Diffusers package, install the pinned submodule:
@@ -74,11 +85,13 @@ This experimental feature leverages `diffusers`'s `transformer.set_attention_bac
   <tr><td rowspan="2">Image(s)-to-Image</td><td><a href="https://huggingface.co/Qwen/Qwen-Image-Edit-2509">Qwen-Image-Edit-2509</a></td><td>20B</td><td>qwen-image-edit-plus</td></tr>
   <tr><td><a href="https://huggingface.co/Qwen/Qwen-Image-Edit-2511">Qwen-Image-Edit-2511</a></td><td>20B</td><td>qwen-image-edit-plus</td></tr>
 
-  <tr><td rowspan="9">Text-to-Image & Image(s)-to-Image</td><td><a href="https://huggingface.co/black-forest-labs/FLUX.2-dev">FLUX.2-dev</a></td><td>32B</td><td>flux2</td></tr>
+  <tr><td rowspan="11">Text-to-Image & Image(s)-to-Image</td><td><a href="https://huggingface.co/black-forest-labs/FLUX.2-dev">FLUX.2-dev</a></td><td>32B</td><td>flux2</td></tr>
   <tr><td><a href="https://huggingface.co/black-forest-labs/FLUX.2-klein-4B">FLUX.2-klein-4B</a></td><td>4B</td><td>flux2-klein</td></tr>
   <tr><td><a href="https://huggingface.co/black-forest-labs/FLUX.2-klein-9B">FLUX.2-klein-9B</a></td><td>9B</td><td>flux2-klein</td></tr>
   <tr><td><a href="https://huggingface.co/black-forest-labs/FLUX.2-klein-base-4B">FLUX.2-klein-base-4B</a></td><td>4B</td><td>flux2-klein</td></tr>
   <tr><td><a href="https://huggingface.co/black-forest-labs/FLUX.2-klein-base-9B">FLUX.2-klein-base-9B</a></td><td>9B</td><td>flux2-klein</td></tr>
+  <tr><td><a href="https://huggingface.co/black-forest-labs/flux-3-image-klein-base">flux-3-image-klein-base</a></td><td>5.5B</td><td>flux3</td></tr>
+  <tr><td><a href="https://huggingface.co/black-forest-labs/flux-3-image-klein">flux-3-image-klein</a></td><td>5.5B</td><td>flux3</td></tr>
   <tr><td><a href="https://huggingface.co/ByteDance-Seed/BAGEL-7B-MoT">BAGEL-7B-MoT</a></td><td>14B</td><td>bagel</td></tr>
   <tr><td><a href="https://huggingface.co/sensenova/SenseNova-U1-8B-MoT">SenseNova-U1 1.0</a></td><td>16B</td><td>sensenova</td></tr>
   <tr><td><a href="https://huggingface.co/sensenova/SenseNova-U1.5-8B-MoT">SenseNova-U1 1.5</a></td><td>16B</td><td>sensenova</td></tr>
@@ -107,7 +120,7 @@ This experimental feature leverages `diffusers`'s `transformer.set_attention_bac
 > To support new models, see [Guidance/New Model](guidance/new_model.md).
 
 > **Offline output support:** SFT and offline DPO currently support `sd3-5`, `flux1`,
-> `flux1-kontext`, `flux2`, `flux2-klein`, `qwen-image`, `qwen-image-edit-plus`,
+> `flux1-kontext`, `flux2`, `flux2-klein`, `flux3`, `qwen-image`, `qwen-image-edit-plus`,
 > `qwen-image-2.1`, `z-image`,
 > `bagel`, `sensenova`, `wan2_t2v`, `wan2_i2v`, `ltx2_t2av`, `ltx2_i2av`, and all
 > MiniMax H3 workflows. Video/audio targets are encoded on demand and are never written to

dataset/offline_smoke/profiles.py

@@ -230,9 +230,11 @@ def _case(
 
 
 _MODEL_TYPES = {
-    "text_to_image": "sd3-5 flux1 flux2 flux2-klein qwen-image qwen-image-2.1 z-image bagel sensenova",
-    "image_to_image": "flux1-kontext flux2 flux2-klein qwen-image-edit-plus bagel sensenova",
-    "multi_image_to_image": "flux2 flux2-klein qwen-image-edit-plus bagel sensenova",
+    "text_to_image": (
+        "sd3-5 flux1 flux2 flux2-klein flux3 qwen-image qwen-image-2.1 z-image bagel sensenova"
+    ),
+    "image_to_image": "flux1-kontext flux2 flux2-klein flux3 qwen-image-edit-plus bagel sensenova",
+    "multi_image_to_image": "flux2 flux2-klein flux3 qwen-image-edit-plus bagel sensenova",
     "text_to_video": "wan2_t2v",
     "first_frame_to_video": "wan2_i2v",
     "first_last_frame_to_video": "wan2_i2v",

examples/README.md

@@ -96,6 +96,14 @@ The same adapter also accepts prompt-only records. It currently requires the bun
 submodule, enforces `per_device_batch_size: 1`, and keeps prefix KV caching disabled for exact
 rollout/replay parity.
 
+## FLUX.3 example
+
+- [`flux3` text-to-image & multi-reference editing GRPO](grpo/lora/flux3/default.yaml)
+
+Supports both `black-forest-labs/flux-3-image-klein-base` (true CFG) and
+`black-forest-labs/flux-3-image-klein` (step-distilled flow map, no CFG) via the bundled
+Diffusers submodule.
+
 ## MiniMax H3 validation status
 
 The T2VA `debug.yaml` recipe is real-weight validated with the 61 GB checkpoint

examples/grpo/lora/flux3/default.yaml

@@ -0,0 +1,113 @@
+# FLUX.3 currently requires the bundled Diffusers submodule:
+#   git submodule update --init
+#   pip install -e ./diffusers
+
+# Environment Configuration
+launcher: "accelerate"
+config_file: config/deepspeed/deepspeed_zero2.yaml  # Switch to config/accelerate_configs/fsdp2.yaml if you have limited GPU memory.
+num_processes: 8
+main_process_port: 29500
+mixed_precision: "bf16"
+
+# Data Configuration
+data:
+  datasets:
+    - name: default
+      dataset_dir: "dataset/sharegpt4o_image_mini"
+      train: { weight: 1 }
+      eval: {}
+  preprocessing_batch_size: 8
+  dataloader_num_workers: 16
+  force_reprocess: true
+  cache_dir: "~/.cache/flow_factory/datasets"
+  sampler_type: "auto"
+  sampler_subgroup_size: null
+
+# Model Configuration
+model:
+  finetune_type: "lora"
+  lora_rank: 64
+  lora_alpha: 128
+  target_modules: "default"
+  model_name_or_path: "black-forest-labs/flux-3-image-klein-base"  # Options: black-forest-labs/flux-3-image-klein-base (CFG), black-forest-labs/flux-3-image-klein (step-distilled flow-map)
+  model_type: "flux3"
+  resume_path: null
+  resume_type: null
+
+log:
+  run_name: null
+  project: "Flow-Factory"
+  logging_backend: "wandb"
+  save_dir: "saves/"
+  save_freq: 20
+  save_model_only: true
+
+# Training Configuration
+train:
+  trainer_type: "grpo"
+  reward_optimization_overlap: false
+  reward_optimization_overlap_mode: "ready"
+  reward_optimization_overlap_poll_interval: 0.05
+  advantage_aggregation: "gdpo"
+  clip_range: 1.0e-4
+  adv_clip_range: 5.0
+  kl_type: "v-based"
+  kl_beta: 0
+  ref_param_device: "cuda"
+
+  # Sampling settings
+  resolution: 384
+  condition_image_size: 384
+  num_inference_steps: 10  # Use 4-16 with guidance_scale: 1.0 for black-forest-labs/flux-3-image-klein
+  guidance_scale: 3.5
+
+  # Batch and sampling
+  per_device_batch_size: 1
+  group_size: 16
+  global_std: false
+  unique_sample_num_per_epoch: 48
+  gradient_step_per_epoch: 2
+  gradient_accumulation_steps: auto
+
+  ema_decay: 0.9
+  ema_update_interval: 4
+  ema_device: "cuda"
+  enable_gradient_checkpointing: true
+  offload_samples_to_cpu: false
+  seed: 42
+
+# Scheduler Configuration
+scheduler:
+  dynamics_type: "Flow-SDE"
+  noise_level: 0.7
+  num_sde_steps: 1
+  sde_steps: [1, 2, 3]
+  seed: 42
+
+# Evaluation Configuration
+eval:
+  resolution: 1024
+  condition_image_size: 1024
+  per_device_batch_size: 1
+  guidance_scale: 3.5
+  num_inference_steps: 28
+  eval_freq: 20
+  seed: 42
+
+# Reward Model Configuration
+rewards:
+  - name: "pick_score"
+    reward_model: "PickScore"
+    weight: 1
+    batch_size: 16
+    device: "cuda"
+    dtype: bfloat16
+
+# Optimizer Configuration
+optimizers:
+  - name: default
+    learning_rate: 3.0e-4
+    weight_decay: 1.0e-4
+    betas: [0.9, 0.999]
+    eps: 1.0e-8
+    max_grad_norm: 1.0