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

miles/algorithms/__init__.py

@@ -0,0 +1,17 @@
+"""Diffusion algorithm plugins.
+
+Currently ships Flow-GRPO only; SFT / AWM / DiffusionNFT land in follow-up PRs.
+"""
+
+from miles.algorithms.base import CollectionSpec, DiffusionAlgorithm, TrainLossContext, TrainSignals
+from miles.algorithms.registry import builtin_algorithm_names, load_algorithm, resolve_algorithm_class_path
+
+__all__ = [
+    "CollectionSpec",
+    "DiffusionAlgorithm",
+    "TrainLossContext",
+    "TrainSignals",
+    "builtin_algorithm_names",
+    "load_algorithm",
+    "resolve_algorithm_class_path",
+]

miles/algorithms/base.py

@@ -0,0 +1,95 @@
+"""Diffusion algorithm plugin interfaces.
+
+Model-family concerns stay in ``TrainPipelineConfig``; algorithm plugins own
+collection contracts, train-example schema, and loss.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Protocol, runtime_checkable
+
+import torch
+
+
+@dataclass(frozen=True)
+class CollectionSpec:
+    """What the collector / rollout engine must provide for this algorithm."""
+
+    mode: str  # "offline" | "online"
+    needs_reward: bool = True
+    needs_trajectory: bool = True
+    needs_logprob: bool = True
+    sampler: str = "sde"  # "sde" | "ode" | "any"
+    return_denoising_env: bool = True
+    sync_weights_to_rollout: bool = True
+
+
+@dataclass
+class TrainLossContext:
+    """Train-side handles passed into ``compute_loss``."""
+
+    models: dict[str, torch.nn.Module]
+    model: torch.nn.Module
+    train_pipeline_config: Any
+    sde_backend: Any | None
+    scheduler: Any | None
+    args: Any
+    forward_dtype: torch.dtype
+    device: torch.device
+
+
+@dataclass
+class LossOutput:
+    """Reserved return type for ``compute_loss`` (loss + metrics).
+
+    **Unused in the Flow-GRPO PR:** ``compute_loss`` still returns a bare
+    ``Tensor`` and appends into the caller's ``log_stats`` (actor-compatible).
+    Kept for a later cleanup that unifies the return shape.
+    """
+
+    loss_sum: torch.Tensor
+    log_stats: dict[str, list[torch.Tensor]] = field(default_factory=dict)
+
+
+@dataclass
+class TrainSignals:
+    """Reward-derived training signals attached to samples before ``build_train_data``.
+
+    Not classification labels: e.g. GRPO advantages or NFT soft +/− weights.
+    """
+
+    raw_rewards: list[float]
+    advantages: list[float] | None = None
+    # Reserved for DiffusionNFT soft positive/negative weights; Flow-GRPO ignores it.
+    nft_signals: list[float] | None = None
+
+
+@runtime_checkable
+class DiffusionAlgorithm(Protocol):
+    name: str
+
+    def validate_args(self, args) -> None: ...
+
+    def collection_spec(self) -> CollectionSpec:
+        """Return acquisition contract; see ``CollectionSpec`` — not fully consumed yet."""
+        ...
+
+    def postprocess_rewards(self, args, samples: list) -> TrainSignals: ...
+
+    def build_train_data(self, args, samples: list, signals: TrainSignals) -> dict[str, Any]: ...
+
+    def validate_train_batch(self, batch: list[dict]) -> list[str]: ...
+
+    def compute_loss(
+        self,
+        ctx: TrainLossContext,
+        batch: list[dict],
+        *,
+        log_stats: dict[str, list[torch.Tensor]],
+        pad_to_len: int | None = None,
+    ) -> torch.Tensor: ...
+
+    def prepare_rollout_data(self, rollout_data: dict, ctx: TrainLossContext) -> None:
+        """Optional hook before the micro-batch loop (e.g. sync scheduler meta)."""
+        ...

miles/algorithms/flow_grpo.py

@@ -0,0 +1,216 @@
+"""Flow-GRPO: reverse-SDE log-prob + PPO-clip."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+from miles.algorithms.base import CollectionSpec, TrainLossContext, TrainSignals
+from miles.algorithms.signals import grpo_group_advantages
+from miles.algorithms.train_forward_utils import (
+    append_rollout_train_abs_diff_stats,
+    compute_noise_pred,
+    prepare_cfg_conds,
+    resolve_cfg_flags,
+    select_model_for_timesteps,
+)
+from miles.utils.train_data_utils import (
+    RolloutTrainDataConverter,
+    scheduler_meta_from_rollout,
+    stack_train_pair_rollout_debug,
+)
+from miles.utils.types import Sample
+
+
+class FlowGRPOAlgorithm:
+    name = "flow_grpo"
+
+    def validate_args(self, args) -> None:
+        kl_beta = float(getattr(args, "diffusion_kl_beta", 0.0) or 0.0)
+        if kl_beta > 0 and not args.use_lora:
+            raise ValueError(
+                "--diffusion-kl-beta currently requires --use-lora so the base model can be used as reference."
+            )
+
+    def collection_spec(self) -> CollectionSpec:
+        return CollectionSpec(
+            mode="online",
+            needs_reward=True,
+            needs_trajectory=True,
+            needs_logprob=True,
+            sampler="sde",
+            return_denoising_env=True,
+            sync_weights_to_rollout=True,
+        )
+
+    def postprocess_rewards(self, args, samples: list[Sample]) -> TrainSignals:
+        return grpo_group_advantages(args, samples)
+
+    def build_train_data(self, args, samples: list[Sample], signals: TrainSignals) -> dict[str, Any]:
+        rewards = signals.advantages if signals.advantages is not None else signals.raw_rewards
+        return RolloutTrainDataConverter().convert_samples(samples, rewards, signals.raw_rewards)
+
+    def validate_train_batch(self, batch: list[dict]) -> list[str]:
+        errors: list[str] = []
+        required = ("latent", "next_latent", "timestep", "next_timestep", "log_prob_old", "advantage", "denoising_env")
+        for i, pair in enumerate(batch):
+            for key in required:
+                if key not in pair:
+                    errors.append(f"batch[{i}] missing {key}")
+        return errors
+
+    def prepare_rollout_data(self, rollout_data: dict, ctx: TrainLossContext) -> None:
+        if ctx.scheduler is None:
+            return
+        num_train_timesteps = ctx.scheduler.config.num_train_timesteps
+        scheduler_timesteps, scheduler_sigmas = scheduler_meta_from_rollout(
+            rollout_data,
+            device=ctx.device,
+            num_train_timesteps=num_train_timesteps,
+        )
+        ctx.scheduler.timesteps = scheduler_timesteps
+        ctx.scheduler.sigmas = scheduler_sigmas
+        ctx.scheduler._step_index = None
+        ctx.scheduler._begin_index = None
+
+    def compute_loss(
+        self,
+        ctx: TrainLossContext,
+        batch: list[dict],
+        *,
+        log_stats: dict[str, list[torch.Tensor]],
+        pad_to_len: int | None = None,
+    ) -> torch.Tensor:
+        """One DiT forward + PPO loss over ``len(batch)`` train pairs. Returns sum of per-pair losses."""
+        if ctx.sde_backend is None:
+            raise RuntimeError("Flow-GRPO requires an SDE step backend")
+
+        args = ctx.args
+        forward_dtype = ctx.forward_dtype
+        train_pipeline_config = ctx.train_pipeline_config
+        device = ctx.device
+
+        use_cfg, guidance_scale, true_cfg_scale = resolve_cfg_flags(args)
+        clip_range = args.diffusion_clip_range
+        noise_level = args.diffusion_noise_level
+        num_train_timesteps = ctx.scheduler.config.num_train_timesteps
+        kl_beta = float(args.diffusion_kl_beta)
+
+        def _stack(key):
+            return torch.stack([pair[key] for pair in batch]).to(device=device, dtype=torch.float32)
+
+        latents_microbatch = _stack("latent")
+        next_latents_microbatch = _stack("next_latent")
+        timesteps_microbatch = _stack("timestep")
+        next_timesteps_microbatch = _stack("next_timestep")
+        log_prob_old_microbatch = _stack("log_prob_old")
+
+        advantage = torch.tensor(
+            [float(pair["advantage"]) for pair in batch],
+            device=device,
+            dtype=torch.float32,
+        )
+        advantage = torch.clamp(advantage, -args.diffusion_adv_clip_max, args.diffusion_adv_clip_max)
+
+        component, model, guidance_scale = select_model_for_timesteps(
+            ctx,
+            timesteps_microbatch,
+            guidance_scale=guidance_scale,
+            num_train_timesteps=num_train_timesteps,
+        )
+
+        if train_pipeline_config.needs_timestep_scaling:
+            timesteps_for_model = timesteps_microbatch / float(num_train_timesteps)
+        else:
+            timesteps_for_model = timesteps_microbatch
+
+        pos_cond, neg_cond, joint_cond, cfg_batching = prepare_cfg_conds(
+            ctx, batch, use_cfg=use_cfg, pad_to_len=pad_to_len
+        )
+
+        latents_input = latents_microbatch.to(forward_dtype)
+        timesteps_input = timesteps_for_model.to(forward_dtype)
+
+        def _pred(disable_adapter: bool = False) -> torch.Tensor:
+            return compute_noise_pred(
+                ctx,
+                model=model,
+                latents_input=latents_input,
+                timesteps_input=timesteps_input,
+                pos_cond=pos_cond,
+                neg_cond=neg_cond,
+                joint_cond=joint_cond,
+                use_cfg=use_cfg,
+                cfg_batching=cfg_batching,
+                guidance_scale=guidance_scale,
+                true_cfg_scale=true_cfg_scale,
+                disable_adapter=disable_adapter,
+            )
+
+        noise_pred_microbatch = _pred()
+
+        _, log_prob_new_microbatch, prev_sample_mean_new, std_dev_t_new = ctx.sde_backend.sde_step_logprob(
+            noise_pred_microbatch.float(),
+            timesteps_microbatch,
+            next_timesteps_microbatch,
+            latents_microbatch.float(),
+            prev_sample=next_latents_microbatch.float(),
+            noise_level=noise_level,
+        )
+
+        log_prob_new = log_prob_new_microbatch
+        log_prob_old = log_prob_old_microbatch
+        ratio = torch.exp(log_prob_new - log_prob_old)
+        unclipped = -advantage * ratio
+        clipped = -advantage * torch.clamp(ratio, 1.0 - clip_range, 1.0 + clip_range)
+        per_pair_loss = torch.maximum(unclipped, clipped)
+        loss_sum = per_pair_loss.sum()
+
+        kl_loss = loss_sum.new_zeros(())
+        if kl_beta > 0:
+            with torch.no_grad():
+                ref_noise_pred_microbatch = _pred(disable_adapter=True)
+                _, _, prev_sample_mean_ref, _ = ctx.sde_backend.sde_step_logprob(
+                    ref_noise_pred_microbatch.float(),
+                    timesteps_microbatch,
+                    next_timesteps_microbatch,
+                    latents_microbatch.float(),
+                    prev_sample=next_latents_microbatch.float(),
+                    noise_level=noise_level,
+                )
+            kl_per_pair = ((prev_sample_mean_new - prev_sample_mean_ref) ** 2).mean(
+                dim=tuple(range(1, prev_sample_mean_new.ndim)),
+                keepdim=True,
+            ) / (2 * std_dev_t_new**2)
+            loss_sum = loss_sum + kl_beta * kl_per_pair.sum()
+            kl_loss = kl_per_pair.mean()
+
+        with torch.no_grad():
+            log_stats["loss"].append((per_pair_loss.mean() + kl_beta * kl_loss).detach())
+            log_stats["policy_loss"].append(per_pair_loss.mean().detach())
+            log_stats["kl_loss"].append(kl_loss.detach())
+            log_stats["loss_abs_mean"].append(per_pair_loss.abs().mean().detach())
+            log_stats["adv_abs_mean"].append(advantage.abs().mean().detach())
+            log_stats["ratio_abs_minus_1"].append((ratio - 1.0).abs().mean().detach())
+            log_stats["approx_kl"].append(0.5 * torch.mean((log_prob_new - log_prob_old) ** 2).detach())
+            log_stats["clipfrac"].append(torch.mean((torch.abs(ratio - 1.0) > clip_range).float()).detach())
+            log_stats["log_prob_new_idx_0"].append(log_prob_new[0].detach())
+            log_stats["log_prob_old_idx_0"].append(log_prob_old[0].detach())
+            log_prob_mean_abs_diff = torch.mean(torch.abs(log_prob_new - log_prob_old)).detach()
+            log_stats["log_prob_mean_abs_diff"].append(log_prob_mean_abs_diff)
+            if len(ctx.models) > 1:
+                log_stats[f"log_prob_mean_abs_diff_{component}"].append(log_prob_mean_abs_diff)
+
+            rollout_model_output = stack_train_pair_rollout_debug(batch, "rollout_step_model_output")
+            if rollout_model_output is not None:
+                mean_abs_diff = append_rollout_train_abs_diff_stats(
+                    log_stats,
+                    "model_output",
+                    noise_pred_microbatch.float(),
+                    rollout_model_output.to(device=device, dtype=torch.float32),
+                )
+                if len(ctx.models) > 1:
+                    log_stats[f"model_output_mean_abs_diff_{component}"].append(mean_abs_diff)
+
+        return loss_sum

miles/algorithms/registry.py

@@ -0,0 +1,44 @@
+"""Resolve ``--diffusion-algorithm`` / ``--diffusion-algorithm-path`` to a class.
+
+PR1 ships Flow-GRPO only; later PRs register additional builtins.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from miles.algorithms.base import DiffusionAlgorithm
+
+_BUILTIN: dict[str, str] = {
+    "flow_grpo": "miles.algorithms.flow_grpo.FlowGRPOAlgorithm",
+}
+
+
+def resolve_algorithm_class_path(args) -> str:
+    """Return the dotted class path for the selected diffusion algorithm."""
+    if getattr(args, "diffusion_algorithm_path", None):
+        return args.diffusion_algorithm_path
+    name = getattr(args, "diffusion_algorithm", None) or "flow_grpo"
+    key = str(name).strip().lower()
+    if key not in _BUILTIN:
+        raise ValueError(
+            f"Unknown --diffusion-algorithm {name!r}; choose one of {sorted(_BUILTIN)} "
+            "or pass --diffusion-algorithm-path. "
+            "SFT/AWM/NFT land in follow-up PRs."
+        )
+    return _BUILTIN[key]
+
+
+def load_algorithm(args) -> DiffusionAlgorithm:
+    from miles.utils.misc import load_function
+
+    path = resolve_algorithm_class_path(args)
+    cls = load_function(path)
+    algo = cls() if isinstance(cls, type) else cls
+    algo.validate_args(args)
+    return algo
+
+
+def builtin_algorithm_names() -> list[str]:
+    return sorted(_BUILTIN)