Customization

Most of miles-diffusion’s behavior can be replaced with user-supplied Python by passing a --*-path flag (loaded via miles.utils.misc.load_function as a dotted import path). This page lists the diffusion-relevant hooks, the signatures they expect, and the defaults they replace — in the same spirit as Miles customization.

At a glance

Stage	Flag	Replaces
Rollout	`--rollout-function-path`	The whole rollout loop
	`--eval-function-path`	The eval rollout
	`--data-source-path`	How prompts / buffer are loaded
	`--custom-generate-function-path`	One microgroup’s generation
	`--diffusion-step-strategy-path`	Which denoising steps enter training
Reward	`--custom-rm-path`	Batched reward computation
	`--custom-reward-post-process-path`	Advantage normalization
Filtering	`--dynamic-sampling-filter-path`	Per-group keep / drop
	`--buffer-filter-path`	Buffer dequeue selection
	`--rollout-sample-filter-path`	Per-sample loss mask
Training	`--custom-expand-samples-to-train-pairs-path`	Trajectory → train pairs
	`--custom-convert-samples-to-train-data-path`	Full sample → train_data
	`--custom-prepare-train-batch-path`	DiT batch preparation
	`--custom-loss-function-path`	Loss formula only
	`--sde-step-backend-path`	Train-side SDE dynamics
Logging	`--custom-rollout-log-function-path`	Train-rollout logging
	`--custom-eval-rollout-log-function-path`	Eval-rollout logging
Model family	`--train-pipeline-config-path`	Family `TrainPipelineConfig`
	`--model-backend-path`	Model load / FSDP backend

--loss-type (policy_loss / nft / sft_loss / custom_loss) is the shortcut that auto-fills several of the train-side paths; see § Training.

Rollout

`--rollout-function-path`

Replace the entire train rollout. Diffusion recipes use miles.rollout.sglang_diffusion_rollout.generate_rollout (not the LLM default).

def generate_rollout(
    args, rollout_id, data_source, evaluation=False
) -> RolloutFnTrainOutput | RolloutFnEvalOutput:
    ...

miles/ray/rollout.py passes the data source positionally through call_rollout_fn. Legacy list/dict returns are still wrapped into the corresponding output dataclass by miles.rollout.base_types.call_rollout_fn.

`--eval-function-path`

Same shape as the rollout function. Defaults to --rollout-function-path when unset.

`--data-source-path`

Class (not a function). Default: miles.rollout.data_source.RolloutDataSourceWithBuffer.

class CustomDataSource:
    def __init__(self, args): ...
    def get_samples(self, num_samples) -> list[list[Sample]]: ...
    def add_samples(self, samples) -> None: ...
    def save(self, rollout_id) -> None: ...
    def load(self, rollout_id=None) -> None: ...

`--custom-generate-function-path`

Replace the inner microgroup generator inside generate_and_rm_microgroup. Default: generate_microgroup.

async def custom_generate(
    args, microgroup: list[Sample], sampling_params: dict, *, evaluation: bool = False
) -> list[Sample]:
    ...

evaluation= is optional — if your signature omits it, the caller skips that kwarg.

`--diffusion-step-strategy-path`

Select which denoising steps contribute SDE log-probs / train pairs. Stock implementations live in miles/rollout/step_strategy_hub.py.

def strategy(args, sample, num_steps, seed) -> tuple[list[int] | None, list[int] | None]:
    # returns (sde_step_indices, return_step_indices); return must be None today
    ...

Stock	Behavior
`sde_window`	Random contiguous window (`--diffusion-num-sde-steps`, `--diffusion-sde-window-range`)
`epoch_global_random_choice`	Per-epoch subset of `--diffusion-sde-candidate-steps`

Details: SDE step backend.

Reward

Built-in scorers (--rm-type pickscore / ocr) are documented in Rewards. The hooks below replace that dispatch entirely.

`--custom-rm-path`

async def custom_rm(args, samples: list[Sample], **kwargs) -> list[float]:
    ...

Wired only through batched_async_rm — implement per-sample routing inside your batched function if needed.

--custom-rm-path my_project.rewards.aesthetic_rm

HTTP / remote scoring: implement a batched custom RM and read args.rm_url (or your own flags). Encode images from sample.generated_output (see _sample_to_rgb_hwc_uint8_frames in miles/rollout/rm_hub/pickscore.py):

import aiohttp
from miles.utils.types import Sample

async def api_rm(args, samples: list[Sample], **kwargs) -> list[float]:
    async with aiohttp.ClientSession() as session:
        rewards = []
        for sample in samples:
            payload = {"prompt": sample.prompt, "image_b64": "<your encoding>"}
            async with session.post(args.rm_url, json=payload) as resp:
                rewards.append((await resp.json())["score"])
        return rewards

--custom-rm-path my_project.rewards.api_rm \
--rm-url http://localhost:8000/score

`--custom-reward-post-process-path`

Replace GRPO advantage normalization in RolloutManager._post_process_rewards.

def post_process(args, samples) -> tuple[list[float], list[float]]:
    # Returns (raw_rewards, normalized_rewards)
    ...

Default behavior: reshape to (-1, n_samples_per_prompt), subtract mean (--globalize-reward-mean for batch-level), optionally divide by std (on by default — --disable-grpo-std-normalization turns it off; --globalize-reward-std switches per-group std to batch-wide).

Filtering

`--dynamic-sampling-filter-path`

Per-group filter after scoring (DAPO-style). Stock: miles.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std.

def filter_function(args, samples: list[Sample], **kwargs):
    # return DynamicFilterOutput(keep=..., reason=...) or a bool
    ...

`--buffer-filter-path`

Select samples when dequeuing from the rollout buffer. Default is pop_first in miles/rollout/data_source.py.

def buffer_filter(
    args, rollout_id: int | None, buffer: list[list[Sample]], num_samples: int
) -> list[list[Sample]]:
    ...

`--rollout-sample-filter-path`

Per-sample, in-place. Set sample.remove_sample = True to exclude a sample from the loss (it still participates in advantage normalization).

def filter_function(args, data: list[list[Sample]]) -> None:
    for group in data:
        for s in group:
            if not_good(s):
                s.remove_sample = True

Training

--loss-type picks the default prepare / formula / expand paths:

`--loss-type`	Expand pairs	Prepare	Loss formula
`policy_loss` (default)	Flow-GRPO built-in	`prepare_flow_grpo_batch`	`flow_grpo_loss_formula`
`nft`	`data_conversion_hub.nft.expand_samples_to_train_pairs`	`loss_hub.nft.prepare_nft_batch`	`nft_loss_formula`
`sft_loss`	(via SFT convert)	required custom / SFT defaults	required custom / SFT defaults
`custom_loss`	your expand (or convert)	your prepare	your formula

DiT forward always stays in the FSDP actor; the loss hook only computes the objective.

`--custom-expand-samples-to-train-pairs-path`

def expand_samples_to_train_pairs(args, samples, rewards, raw_rewards) -> dict:
    ...

Default for Flow-GRPO lives under miles/ray/data_conversion_hub/flow_grpo.py.

`--custom-convert-samples-to-train-data-path`

Replace the entire RolloutManager._convert_samples_to_train_data (including reward post-process + expand). Prefer the narrower expand hook unless you need full control.

def convert_samples_to_train_data(args, samples) -> dict:
    ...

`--custom-prepare-train-batch-path`

def prepare(ctx, batch, *, pad_to_len=None) -> PreparedBatch:
    ...

Builds DiT inputs from train pairs. Defaults: miles.backends.fsdp_utils.loss_hub.flow_grpo.prepare_flow_grpo_batch or the NFT equivalent under loss_hub.nft.

`--custom-loss-function-path`

def loss_formula(ctx, batch, prepared, *, new_pred, ref_pred, metrics, **kwargs):
    ...

`--sde-step-backend-path`

Class implementing SdeStepBackend. Auto-selected from --diffusion-sde-type (sde / ode → DiffusersSdeStepBackend, cps → CpsSdeStepBackend) unless overridden. See SDE step backend.

Logging

def log_rollout_data(rollout_id, args, samples, rollout_extra_metrics, rollout_time) -> bool:
    ...

def log_eval_rollout_data(rollout_id, args, data, extra_metrics) -> bool:
    ...

Return a truthy value to skip the default logging; falsy layers on top.

Model family

`--hf-checkpoint` / `--diffusion-model-family` / `--train-pipeline-config-path`

--hf-checkpoint names the diffusers pipeline for train + rollout. Family is resolved from the checkpoint name unless you pass --diffusion-model-family (e.g. sd3). For an unregistered family, pass --train-pipeline-config-path to a TrainPipelineConfig subclass instead.

`--model-backend-path`

Class for loading components / FSDP / scheduler. Default comes from the family config (usually a DiffusersModelBackend).

Worked example

Custom reward + post-process on top of the stock SD3 Flow-GRPO recipe:

python3 scripts/run_diffusion_grpo_sd3_ocr_sglang.py \
  --cuda-visible-devices 6,7 \
  --extra-args "
    --custom-rm-path my_project.rewards.api_rm
    --rm-url http://localhost:8000/score
    --custom-reward-post-process-path my_project.rewards.normalize
  "

(--extra-args forwarding depends on the launch script; you can also splice the flags into a forked recipe.)

Pairs well with

Rewards — built-in PickScore / OCR and prompt JSONL.
SDE step backend — step strategies and SDE kernels.
LoRA weight sync — IPC merge path used by most recipes.
SD3 model guide — end-to-end recipe flags.

Getting Started

User Guide

Models

Advanced

Developer

Customization

At a glance

Rollout

`--rollout-function-path`

`--eval-function-path`

`--data-source-path`

`--custom-generate-function-path`

`--diffusion-step-strategy-path`

Reward

`--custom-rm-path`

`--custom-reward-post-process-path`

Filtering

`--dynamic-sampling-filter-path`

`--buffer-filter-path`

`--rollout-sample-filter-path`

Training

`--custom-expand-samples-to-train-pairs-path`

`--custom-convert-samples-to-train-data-path`

`--custom-prepare-train-batch-path`

`--custom-loss-function-path`

`--sde-step-backend-path`

Logging

Model family

`--hf-checkpoint` / `--diffusion-model-family` / `--train-pipeline-config-path`

`--model-backend-path`

Worked example

Pairs well with

​At a glance

​Rollout

​--rollout-function-path

​--eval-function-path

​--data-source-path

​--custom-generate-function-path

​--diffusion-step-strategy-path

​Reward

​--custom-rm-path

​--custom-reward-post-process-path

​Filtering

​--dynamic-sampling-filter-path

​--buffer-filter-path

​--rollout-sample-filter-path

​Training

​--custom-expand-samples-to-train-pairs-path

​--custom-convert-samples-to-train-data-path

​--custom-prepare-train-batch-path

​--custom-loss-function-path

​--sde-step-backend-path

​Logging

​Model family

​--hf-checkpoint / --diffusion-model-family / --train-pipeline-config-path

​--model-backend-path

​Worked example

​Pairs well with

At a glance

Rollout

`--rollout-function-path`

`--eval-function-path`

`--data-source-path`

`--custom-generate-function-path`

`--diffusion-step-strategy-path`

Reward

`--custom-rm-path`

`--custom-reward-post-process-path`

Filtering

`--dynamic-sampling-filter-path`

`--buffer-filter-path`

`--rollout-sample-filter-path`

Training

`--custom-expand-samples-to-train-pairs-path`

`--custom-convert-samples-to-train-data-path`

`--custom-prepare-train-batch-path`

`--custom-loss-function-path`

`--sde-step-backend-path`

Logging

Model family

`--hf-checkpoint` / `--diffusion-model-family` / `--train-pipeline-config-path`

`--model-backend-path`

Worked example

Pairs well with