Workflow Guidance

Table of Contents

• Overview
• Stage 1: Data Preprocessing
• Stage 2: K-Repeat Sampling
• Stage 3: Trajectory Generation
• Stage 4: Reward Computation
• Stage 5: Advantage Computation
• Stage 6: Policy Optimization
• Putting It All Together

Overview

Flow-Factory follows an online RL training paradigm for diffusion/flow-matching models. Each epoch executes a six-stage pipeline:

┌─────────────────────────────────────────────────────────────────────────────────┐
│                          Flow-Factory Training Epoch                            │
│                                                                                 │
│  ┌─────────────┐    ┌──────────┐    ┌──────────────┐    ┌───────────────┐       │
│  │    Data     │    │ K-Repeat │    │  Trajectory  │    │    Reward     │       │
│  │Preprocessing│───►│ Sampling │───►│  Generation  │───►│ Computation   │       │
│  │  (offline)  │    │          │    │  (Adapter)   │    │               │       │
│  └─────────────┘    └─────▲────┘    └──────────────┘    └───────┬───────┘       │
│                           │                                     │               │
│                           │    ┌──────────────┐    ┌────────────▼─-─┐           │
│                           │    │   Policy     │    │   Advantage    │           │
│                           └────│ Optimization │◄───│  Computation   │           │
│                                └──────────────┘    └────────────────┘           │
└─────────────────────────────────────────────────────────────────────────────────┘

The high-level training loop (shared by all algorithms) is defined in each trainer's start() method:

# src/flow_factory/trainers/grpo.py — GRPOTrainer.start()
def start(self):
    while self.should_continue_training():
        # Checkpoint & Evaluation (omitted for brevity)
        samples = self.sample()            # Stages 2 + 3
        self.prepare_feedback(samples)     # Stages 4 + 5 (rewards + advantages)
        self.optimize(samples)             # Stage 6 (DPO: pair formation + loss here)
        self.epoch += 1

Note: Stage 1 (preprocessing) runs once before training begins and is cached to disk. Stages 2–6 repeat every epoch. The three methods above map directly to those stages: sample → trajectory rollouts; prepare_feedback → finalize rewards from the buffer and compute advantages; optimize → policy update (DPO additionally forms chosen/rejected pairs at the start of optimize before the loss).

Stage 1: Data Preprocessing

Goal: Encode raw text prompts (and optional images / videos / audio files) into model-ready tensor representations before training begins, eliminating redundant computation during the RL loop and enabling components offloading such as text-encoder, image-encoder, and audio-encoder (when applicable).

Input / Output

	Description
Input	Raw dataset: train.jsonl or train.txt containing prompts, optional image / video / audio paths
Output	Cached HuggingFace Dataset on disk with pre-encoded tensors (prompt_embeds, prompt_ids, pooled_prompt_embeds, image_latents, etc.)

How It Works

Each model adapter exposes a preprocess_func that encodes raw inputs into tensors. The GeneralDataset class orchestrates this via HuggingFace's .map() with automatic caching:

# src/flow_factory/data_utils/dataset.py — GeneralDataset._preprocess_batch()
def _preprocess_batch(self, batch, image_dir, video_dir, audio_dir):
    # 1. Prepare text prompts
    prompt = batch["prompt"]
    # 2. Load images from disk (if applicable)
    # 3. Load videos from disk (if applicable)
    # 4. Load audio files from disk (if applicable, via utils.audio.load_audio)
    # 5. Call model-specific preprocess function
    preprocess_res = self._preprocess_func(**filtered_args)
    # 6. Move tensors to CPU for caching
    # 7. Return batch dict with encoded tensors + metadata

The preprocess function is model-specific. For example, Flux.2 encodes prompts via its text encoder and images via its VAE:

# src/flow_factory/models/flux/flux2.py — Flux2Adapter.preprocess_func()
def preprocess_func(self, prompt, images, ...):
    batch = self.encode_prompt(prompt=prompt, ...)       # → prompt_embeds, prompt_ids
    if has_images:
        batch.update(self.encode_image(images=images, ...))  # → image_latents, image_ids
    return batch

Audio is symmetric: audio_dir is the third optional input handled by _preprocess_batch, parallel to image_dir / video_dir. Audio-aware adapters (e.g. the LTX-2 audio-video adapter) override encode_audio to consume the loaded audios batch; text/image/video-only adapters inherit the no-op BaseAdapter.encode_audio and ignore the column entirely.

Key Points

• Distributed preprocessing: When running on multiple GPUs, each rank processes a shard of the dataset independently. The orchestrator (loader._create_or_load_dataset) routes each rank's Dataset.map output directly to its final per-rank Arrow file via cache_file_name=, so a shard is written to disk exactly once. After all ranks finish, the consolidator (local-main for preprocess_parallelism="local", global rank 0 for "global") writes only state.json and dataset_info.json referencing the existing per-rank files and atomically renames .tmp → final cache directory — no row data is re-serialized.
• Cache layout: The merged cache directory looks like {cache_dir}/{fingerprint}/_parts/rank_{i:05d}_of_{N:05d}/cache-{fingerprint}_shard{i}of{N-1}.arrow, plus the top-level state.json and dataset_info.json. While preprocessing is in flight, the same content lives under {cache_dir}/{fingerprint}.tmp/, with a _build_meta.json sentinel that records num_shards so a subsequent run with the same num_shards can resume from any per-rank Arrow files that were already written before a crash, while a different num_shards triggers a clean wipe.
• No HF default-cache copy: Because each map() call sets cache_file_name, HuggingFace does not also write a duplicate cache-*.arrow under ~/.cache/huggingface/datasets/....
• Intelligent caching: A hash fingerprint of (dataset, split, max_dataset_size, preprocess_func source, preprocess_kwargs, extra_hash_strs) (the last includes model_type and model_name_or_path) determines the cache path. Subsequent runs that match the fingerprint take the fast path without any Dataset.map invocation.
• Component offloading: Text encoders and VAEs are loaded for preprocessing, then offloaded before the training loop to free VRAM for the denoising model.

Configuration

data:
  dataset: "path/to/dataset"
  enable_preprocess: true          # Enable offline preprocessing
  force_reprocess: false           # Force re-encoding even if cache exists; essential if code is modified without changing config
  preprocessing_batch_size: 16     # Batch size for encoding
  cache_dir: "~/.cache/flow_factory/datasets"
  preprocess_parallelism: "local"  # "local" = per-node parallelism (no shared FS required); "global" = cross-node (shared FS required)

Stage 2: K-Repeat Sampling

Goal: Construct batches where each unique prompt appears exactly $K$ times (group_size), enabling group-relative advantage computation.

Input / Output

	Description
Input	Preprocessed dataset of $N$ samples
Output	Batches of encoded prompts, where each prompt is repeated $K$ times across the distributed cluster

How It Works

The DistributedKRepeatSampler handles this:

# src/flow_factory/data_utils/sampler.py — DistributedKRepeatSampler.__iter__()
def __iter__(self):
    while True:
        g = torch.Generator()
        g.manual_seed(self.seed + self.epoch)
        # 1. Randomly select M unique prompts
        indices = torch.randperm(len(self.dataset), generator=g)[:self.m].tolist()
        # 2. Repeat each prompt K times → M*K total samples
        repeated = [idx for idx in indices for _ in range(self.k)]
        # 3. Distribute evenly across all GPUs
        per_rank = chunk(repeated, self.num_replicas)[self.rank]
        # 4. Yield batches of size `per_device_batch_size`
        yield from chunk(per_rank, self.batch_size)

Key Points

• Deterministic seeding: All ranks share the same seed + epoch generator, ensuring identical permutation and K-repeat ordering — no explicit cross-rank communication needed.
• Automatic alignment: The sampler adjusts unique_sample_num upward to ensure M * K is evenly divisible by batch_size * num_replicas.
• Group identification: Each sample carries a unique_id (hash of prompt + conditions). During advantage computation, samples are grouped by this ID across all ranks.

Configuration

train:
  per_device_batch_size: 2       # Batch size per GPU
  group_size: 4                  # K — repetitions per prompt
  unique_sample_num_per_epoch: 64  # M — unique prompts per epoch

Effective samples per epoch = $M \times K$. For example, with M=64, K=4, each epoch generates 256 samples across the cluster.

Stage 3: Trajectory Generation

Goal: Run the denoising model to generate images/videos from noise, collecting the necessary denoising trajectory (latents and log-probabilities at each timestep).

Input / Output

	Description
Input	Batched raw input (prompt, images) or encoded tensors (prompt_embeds, image_latents) from the dataloader.
Output	List[BaseSample] — each sample contains: generated image/video, denoising trajectory (all_latents), log-probabilities (log_probs), timestep schedule, and prompt info

How It Works

The trainer's sample() method switches the adapter to rollout mode and runs inference:

# src/flow_factory/trainers/grpo.py — GRPOTrainer.sample()
def sample(self) -> List[BaseSample]:
    trajectory_indices = compute_trajectory_indices(
        train_timestep_indices=self.adapter.scheduler.train_timesteps,
        num_inference_steps=self.training_args.num_inference_steps,
    )
    # generate_samples() (BaseTrainer) switches the adapter to rollout mode,
    # loops the dataloader, runs adapter.inference() under no_grad + autocast,
    # buffers rewards (reward_buffer), and returns the collected samples.
    return self.generate_samples(
        reward_buffer=self.reward_buffer,
        compute_log_prob=True,
        trajectory_indices=trajectory_indices,
    )

Inside adapter.inference(), the model runs a multi-step denoising loop (SDE or ODE), collecting latents and computing log-probabilities at each step. The result is packaged into BaseSample dataclass instances:

# Example: src/flow_factory/models/flux/flux1.py — Inference result
BaseSample(
    timesteps=timesteps,                # (T+1,) schedule
    all_latents=stacked_latents,        # (num_stored, seq_len, C) — selectively stored
    log_probs=stacked_log_probs,        # (num_stored,) — per-step log π(a|s)
    latent_index_map=latent_index_map,  # (T+1,) maps step → storage index
    log_prob_index_map=log_prob_index_map,
    image=decoded_image,                # (C, H, W) tensor
    prompt=prompt_text,
    prompt_embeds=prompt_embeds,
    ...
)

Algorithm-Specific Differences

Algorithm	compute_log_prob	trajectory_indices	Notes
GRPO	True	Only train timesteps	Needs log-prob for policy ratio; selective storage saves memory.
DiffusionNFT	False	[-1] (final only)	Only needs final clean latent $x_1$; log-prob not required
AWM	False	[-1] (final only)	Same as NFT; log-prob computed later during optimization
DGPO	False	[-1] (final only)	Same trajectory policy as NFT/AWM; optimization uses fresh TimeSampler timesteps
CRD	False	[-1] (final only)	Same trajectory policy as NFT/AWM; reward distillation against CFG-guided teacher reference

Key Points

• Selective trajectory recording: trajectory_indices controls which denoising steps are stored. For GRPO, only steps corresponding to train_timesteps are kept to reduce memory.
• SDE dynamics for exploration: GRPO injects noise during sampling via SDE formulation, enabling the log-probability computation required for policy gradients. NFT, AWM, DGPO, and CRD use decoupled sampling (typically ODE) with compute_log_prob=False.
• Off-policy sampling: NFT optionally uses EMA parameters for sampling (off_policy: true), while the current policy is optimized — stabilizing training.

Stage 4: Reward Computation

Goal: Score each generated sample using one or more reward models.

Input / Output

	Description
Input	List[BaseSample] with generated images/videos and prompts
Output	Dict[str, Tensor] — reward name → per-sample scores (aligned with local samples)

How It Works

The RewardProcessor handles batched, distributed reward computation:

# src/flow_factory/rewards/reward_processor.py — RewardProcessor.compute_rewards()
def compute_rewards(self, samples, store_to_samples=True, epoch=0, split='all'):
    results = {}
    # Pointwise rewards: local computation per rank
    if self._pointwise_models:
        results.update(self._compute_pointwise_rewards(samples, epoch))
    # Groupwise rewards: gather → compute → scatter
    if self._groupwise_models:
        results.update(self._compute_groupwise_rewards(samples, epoch))
    # Store rewards in each sample's extra_kwargs
    if store_to_samples:
        for i, sample in enumerate(samples):
            sample.extra_kwargs['rewards'] = {k: v[i] for k, v in results.items()}
    return results

Key Points

• Pointwise vs Groupwise: Pointwise models (e.g., PickScore, CLIP) compute rewards independently per sample — no cross-rank communication needed. Groupwise models (e.g., ranking-based) require gathering all group members first.
• Automatic deduplication: If multiple reward entries share the same model config, they reuse a single model instance.
• Flexible inputs: Reward models declare required_fields (e.g., ("prompt", "image")) and optionally receive raw tensors (use_tensor_inputs=True) or PIL images.
• Remote reward servers: For reward models with incompatible dependencies, Flow-Factory supports HTTP-based reward computation in isolated environments.

Configuration

rewards:
  - name: "aesthetic"
    reward_model: "PickScore"
    weight: 1.0
    batch_size: 16
  - name: "text_align"
    reward_model: "CLIP"
    weight: 0.5
    batch_size: 32

See Reward Guidance for detailed reward model configuration.

Stage 5: Advantage Computation

Goal: Convert raw rewards into normalized, group-relative advantages that serve as the optimization signal.

Input / Output

	Description
Input	Per-sample rewards (Dict[str, Tensor]) and sample list with unique_id
Output	Per-sample advantage scalar stored in sample.extra_kwargs['advantage']

How It Works

# src/flow_factory/trainers/grpo.py — GRPOTrainer.compute_advantages()
def compute_advantages(self, samples, rewards, store_to_samples=True, aggregation_func=None):
    # Thin wrapper: resolve the aggregation strategy, then delegate to
    # AdvantageProcessor (advantage/advantage_processor.py). The processor is
    # communication-aware and auto-selects the gather-vs-local path; it performs
    # the gather -> weighted-aggregate -> group-by-unique_id -> normalize ->
    # scatter sequence summarized below.
    aggregation_func = aggregation_func or self.training_args.advantage_aggregation
    return self.advantage_processor.compute_advantages(
        samples=samples,
        rewards=rewards,
        store_to_samples=store_to_samples,
        aggregation_func=aggregation_func,
    )

Aggregation Strategies

Strategy	Formula	Use Case
sum	$A = \text{normalize}(\sum_i w_i \cdot r_i)$	Default GRPO: advantage of weighted reward sum
gdpo	$A = \text{BN}(\sum_i w_i \cdot A_i)$	Per-reward normalization first, then combine

Key Points

• Cross-rank synchronization: Advantages are computed globally — rewards from all ranks are gathered, normalized, then scattered back. This ensures consistent group-level statistics.
• Group-relative normalization: Within each group (same prompt), rewards are zero-centered and variance-normalized. This makes the advantage signal invariant to absolute reward scale.
• Batch normalization (GDPO): For multi-reward scenarios, GDPO normalizes each reward independently before combining, preventing one reward from dominating.

Configuration

train:
  advantage_aggregation: 'sum'    # Options: 'sum', 'gdpo'
  global_std: false               # Use global std instead of per-group std
  adv_clip_range: [-5.0, 5.0]    # Clip advantages to prevent outliers

Stage 6: Policy Optimization

Goal: Update the denoising model's parameters using the computed advantages and PPO-style clipped policy gradient.

Input / Output

	Description
Input	List[BaseSample] with advantages, trajectories, and log-probs stored
Output	Updated model parameters; logged loss metrics

How It Works (GRPO)

Stages 4–5 run in prepare_feedback() (reward buffer finalize, then AdvantageProcessor). Stage 6 is optimize() only:

# Stages 4–5 — src/flow_factory/trainers/grpo.py — GRPOTrainer.prepare_feedback()
def prepare_feedback(self, samples):
    rewards = self.reward_buffer.finalize(store_to_samples=True, split='all')
    self.compute_advantages(samples, rewards, store_to_samples=True)
    # ... log advantage metrics ...

# Stage 6 — GRPOTrainer.optimize()
def optimize(self, samples):
    for inner_epoch in range(num_inner_epochs):
        # Shuffle and re-batch
        shuffled = permute(samples)
        batches = [BaseSample.stack(chunk) for chunk in chunks(shuffled)]

        self.adapter.train()
        for batch in batches:
            # Iterate through train timesteps
            for timestep_index in scheduler.train_timesteps:
                with accelerator.accumulate(*trainable_components):
                    # 1. Get old log-prob from trajectory
                    old_log_prob = batch['log_probs'][log_prob_idx]
                    # 2. Forward pass → new log-prob
                    output = self.adapter.forward(latents=x_t, t=t, ...)
                    # 3. PPO-style clipped loss
                    ratio = exp(output.log_prob - old_log_prob)
                    unclipped = -adv * ratio
                    clipped   = -adv * clamp(ratio, 1-ε, 1+ε)
                    loss = mean(max(unclipped, clipped))
                    # 4. Optional KL regularization
                    if enable_kl_loss:
                        loss += kl_beta * KL(current || reference)
                    # 5. Backward + optimizer step
                    accelerator.backward(loss)
                    optimizer.step()

shuffle_samples and on-policy ratio: the optimize loop reorders samples each inner epoch (train.shuffle_samples: true, the default). For adapters whose batched forward() is pack-composition-dependent (e.g. Bagel NaViT packing), this makes a training micro-batch pack a different sample set than its rollout pack, so the on-policy ratio != 1. Set train.shuffle_samples: false for such adapters (with matched sampling/training per_device_batch_size) so each micro-batch reproduces its rollout pack. See the train-inference consistency topic doc.

Algorithm-Specific Optimization

Algorithm	Optimization Strategy
GRPO	Iterates over stored trajectory timesteps; computes ratio from old/new log-probs; PPO clipping
GRPO-Guard	Same as GRPO but with timestep-dependent loss reweighting to mitigate ratio bias
DiffusionNFT	Samples fresh timesteps; interpolates $x_t = (1-t)x_1 + t\epsilon$; contrastive objective with normalized rewards
AWM	Samples fresh timesteps; weights velocity matching loss by advantage; PPO clipping + EMA-KL regularization
DGPO	Samples fresh timesteps via TimeSampler; applies group-level preference objective with optional PPO clipping and EMA-reference KL
CRD	Samples fresh timesteps; reward distillation against CFG-guided teacher with adaptive KL; old/sampling model snapshots and centered advantages
DPO	Preference loss on chosen/rejected pairs; pairs formed at the start of optimize after advantages

Key Points

• Inner epochs: Samples can be reused for multiple optimization passes (num_inner_epochs), amortizing the cost of sampling.
• Gradient accumulation: The accelerator.accumulate() context handles gradient accumulation across timesteps and micro-batches, with optimizer steps only at sync boundaries.
• KL regularization: Optional penalty keeping the policy close to a reference model (or EMA model for AWM), preventing reward hacking.
• Per-timestep iteration: GRPO iterates over each stored trajectory timestep, computing loss at each. NFT, AWM, DGPO, and CRD sample fresh timesteps independently of the sampling trajectory.

Putting It All Together

A complete epoch with GRPO on a 8×GPU cluster:

Epoch N
├── DataLoader (DistributedKRepeatSampler)
│   └── Select 64 unique prompts × 4 repeats = 256 samples
│       → 32 samples per GPU (256 / 8)
│       → 16 batches per GPU (32 / batch_size=2)
│
├── Sampling (torch.no_grad)
│   └── For each batch: adapter.inference(compute_log_prob=True)
│       → 32 BaseSample per GPU, each with trajectory + log-probs
│
├── prepare_feedback(samples)
│   ├── Reward computation: RewardProcessor / buffer finalize → Dict[str, Tensor(32,)] per GPU
│   └── Advantage computation: gather → group by unique_id → normalize → scatter
│
└── optimize(samples) — Stage 6 only (num_inner_epochs × batches × timesteps)
    ├── Shuffle 32 samples → re-batch
    ├── For each batch, for each timestep:
    │   ├── Forward pass → new log-prob
    │   ├── PPO-clipped loss with advantage
    │   ├── + Optional KL penalty
    │   └── Backward + gradient accumulation
    └── Optimizer step at sync boundaries

DPO: form chosen/rejected pairs at the start of optimize() (after advantages exist), then run the preference loss; there is no pair formation in prepare_feedback().