refactor: extract Flow-GRPO into a DiffusionAlgorithm plugin (prep for NFT / SFT)

[niehen6174]

Motivation

We want to support Flow-GRPO, DiffusionNFT, and SFT in one stack. These are not “different loss names on the same train pairs” — they diverge by process direction, which changes collection, train signals, example schema, and loss:

	Flow-GRPO	DiffusionNFT	SFT
Process family	Reverse PG (SDE transitions)	Forward flow matching	Forward flow matching
Data source	Online rollout	Online rollout	Offline dataset
Need reward	Yes → group advantage	Yes → soft +/− signals	Usually no RL reward
Need full traj + logπ	Yes	No (clean x0 is enough)	No
Train example	(x_t, x_{t+Δ}, logπ_old, A)	(x0, nft_signals)	(x0, …)
Loss	PPO-clip on importance ratio	FM with implicit +/− branches	Plain FM MSE

flowchart TB
  subgraph grpo ["Flow-GRPO (reverse PG)"]
    direction LR
    G1["Collect: online SDE<br/>full trajectory + logπ"] --> G2["Signals: group advantage A"]
    G2 --> G3["Example: x_t, x_tΔ, logπ_old, A"]
    G3 --> G4["Loss: PPO-clip on logπ ratio"]
  end

  subgraph nft ["DiffusionNFT (forward matching)"]
    direction LR
    N1["Collect: online<br/>clean x0 + reward<br/>"] --> N2["Signals: soft +/− nft_signals"]
    N2 --> N3["Example: x0 + nft_signals"]
    N3 --> N4["Loss: FM with +/− branches"]
  end

  subgraph sft ["SFT (forward matching)"]
    direction LR
    S1["Collect: offline dataset<br/>x0 + prompt"] --> S2["Signals: none / trivial"]
    S2 --> S3["Example: x0"]
    S3 --> S4["Loss: plain FM MSE"]
  end

  grpo ~~~ nft
  nft ~~~ sft

  style G1 fill:#E8F1FF,stroke:#2F6FED
  style G4 fill:#E8F1FF,stroke:#2F6FED
  style N1 fill:#EAF8EF,stroke:#1B7F4E
  style N4 fill:#EAF8EF,stroke:#1B7F4E
  style S1 fill:#FFF6E5,stroke:#C47D00
  style S4 fill:#FFF6E5,stroke:#C47D00

Loading

Today those Flow-GRPO assumptions are baked into the actor / converter (log_prob_old, step pairs, PPO). Adding NFT/SFT with if algorithm == … would keep growing that surface. The right cut is to keep infra and model family shared, and put algorithm decisions behind a plugin:

DiffusionAlgorithm
  collection_spec()          # what to collect (traj / logπ / online|offline)
  postprocess_rewards(...)   # reward → TrainSignals
  build_train_data(...)      # samples + signals → train examples
  compute_loss(...)          # algorithm objective

flowchart TB
  subgraph shared ["Shared (stable)"]
    Loop["train_diffusion loop"]
    RM["RolloutManager / Collector"]
    Actor["FSDP actor: microbatch / optim"]
    Pipe["TrainPipelineConfig: DiT / CFG"]
  end

  subgraph plugin ["DiffusionAlgorithm (swappable)"]
    Spec["collection_spec"]
    Rew["postprocess_rewards"]
    Build["build_train_data"]
    Loss["compute_loss"]
  end

  subgraph impl ["Implementations"]
    GRPO["FlowGRPO — this PR"]
    NFT["DiffusionNFT — later"]
    SFTAlg["SFT — later"]
  end

  Loop --> RM --> Spec --> Rew --> Build --> Actor --> Loss --> Pipe

  GRPO -.-> Spec
  GRPO -.-> Rew
  GRPO -.-> Build
  GRPO -.-> Loss
  NFT -.-> Spec
  NFT -.-> Rew
  NFT -.-> Build
  NFT -.-> Loss
  SFTAlg -.-> Spec
  SFTAlg -.-> Rew
  SFTAlg -.-> Build
  SFTAlg -.-> Loss

  style plugin fill:#F3F0FF,stroke:#6B4FBB
  style GRPO fill:#E8F1FF,stroke:#2F6FED
  style NFT fill:#EAF8EF,stroke:#1B7F4E
  style SFTAlg fill:#FFF6E5,stroke:#C47D00

Loading

TrainSignals = reward-derived scalars on each sample before building train examples (GRPO advantages, future NFT soft weights) — not classification labels.

---

Modifications (this PR)

Area	Change
miles/algorithms/	New plugin package: base (DiffusionAlgorithm, TrainSignals, CollectionSpec), flow_grpo, signals, train_forward_utils, registry
miles/backends/fsdp_utils/actor.py	Thin train loop: validate_train_batch → prepare_rollout_data → compute_loss (PPO moved out of actor)
miles/ray/rollout.py	Reward postprocess + train-data build delegated to algorithm
miles/utils/arguments.py	--diffusion-algorithm (default flow_grpo) + optional --diffusion-algorithm-path
miles/utils/diffusion_protocol.py	Comments clarifying GRPO keys are the default algorithm contract
tests/fast/algorithms/	Registry load + GRPO advantage normalization smoke tests

Intentionally not changed

• train_diffusion outer loop shape
• TrainPipelineConfig / model-family forward
• Rollout sampling still follows today’s Flow-GRPO SDE path (collection_spec is declared but not yet driving sampling)
• No NFT / SFT / offline collector code

---

Follow-up plan

flowchart LR
  PR1["This PR: Flow-GRPO plugin<br/>behavior unchanged"] --> PR2["DiffusionNFT<br/>forward FM + signals"]
  PR2 --> PR3["SFT<br/>offline collector + FM"]

  style PR1 fill:#E8F1FF,stroke:#2F6FED
  style PR2 fill:#EAF8EF,stroke:#1B7F4E
  style PR3 fill:#FFF6E5,stroke:#C47D00

Loading

1. DiffusionNFT — forward-matching kernel + nft_signals; relax collection (no traj/logπ); still online + reward.
2. SFT — offline collector / dataset path; plain FM MSE; first time train_diffusion collect side may need to be pluggable.
3. (Optional) AWM — mostly another forward-matching plugin (A * FM), small incremental PR after NFT.
4. Wire collection_spec into rollout sampling params so NFT/SFT do not pay for unused SDE trajectory.

---

Consistency verification

Refactor intent: Flow-GRPO numerically / behaviorally unchanged vs pre-refactor baseline on the main model families.

Model family	Verification so far	Result
SD3.5	Smoke test (short Flow-GRPO run)	Pass
Qwen-Image	Smoke test (short Flow-GRPO run)	Pass
LTX	Smoke test (short Flow-GRPO run)	Pass

---

[niehen6174]

This PR will be closed and continued in #63 .