> ## Documentation Index > Fetch the complete documentation index at: https://miles.radixark.com/docs/llms.txt > Use this file to discover all available pages before exploring further. # CLI Reference > Every flag train_diffusion.py accepts, grouped by subsystem. miles-diffusion is configured entirely through flags on `train_diffusion.py`, coming from three places: | Source | Where | Notes | | ---------------------- | ------------------------------------------------------ | ----------------------------------------------------------------------------------------------------------------------------------------------- | | Training-backend flags | `FSDPArgs` in `miles/backends/fsdp_utils/arguments.py` | Every dataclass field becomes `--field-name` automatically | | miles flags | `miles/utils/arguments.py` | Added on top as an `extra_args_provider` | | sglang-d passthrough | The rollout engine's own CLI | Re-registered with a `--sglang-` prefix; a short skip list covers what miles sets itself (`model_path`, ports, `base_gpu_id`, `random_seed`, …) | `--custom-config-path ` loads YAML keys into the namespace after parsing. A key that collides with an existing argument **overrides it** (with a logged warning) — the YAML wins over the command line. Reading the prefixes: | Prefix | Marks | | ---------------------- | ------------------------------------------------------------------------------------------------------------------------------------------- | | `diffusion-` | The modality — denoising, SDE, CFG, latents, frames. Generic ML/RL concepts (clipping, KL, EMA, LoRA, batching) do not take it | | `fsdp-` | The training side — compare `--fsdp-flow-shift` (training-side sigma grid) with `--diffusion-flow-shift` (the engine's generation schedule) | | `rollout-` / `sglang-` | The engine side | The prefix does not tell you which argument group a flag lives in — groups follow concern, not name. `python3 train_diffusion.py --help` is always the ground truth. *** ## Essentials ### The one required flag | Flag | What | | ----------------- | ----------------------------------------------------------------------- | | `--hf-checkpoint` | The diffusers pipeline to train, as an HF repo id or a local directory. | One value serves three readers: training loads components and scheduler from it, the sglang-d engine serves it, and the **model family** is matched from its name. Add `--diffusion-model-family` when the name carries no family hint — which local weights usually do not. ### Cluster topology | Flag | Default | What | | ------------------------------- | ------- | ----------------------------------------------------------------------------------------------- | | `--actor-num-nodes` | `1` | Nodes for the training actor. | | `--actor-num-gpus-per-node` | `8` | GPUs per actor node. Train world size is the product. | | `--rollout-num-gpus` | – | GPUs for rollout-side work. Forced to the train world size under `--colocate`. | | `--rollout-num-gpus-per-engine` | `1` | GPUs per sglang-d engine (its TP × SP). | | `--num-gpus-per-node` | `8` | Total GPUs the job may use per node. **Set this when using fewer than 8.** | | `--colocate` | off | Required: CUDA IPC weight sync currently supports only colocated trainer and rollout execution. | ### Batch sizing | Flag | Default | What | | -------------------------------------------------------- | ------------ | ----------------------------------------------------------------------------------------- | | `--rollout-batch-size` | unset | Prompts per rollout. No default — every recipe sets it. | | `--n-samples-per-prompt` | `1` | Samples per prompt (GRPO group size). | | `--global-batch-size` | derived | **Samples** per optimizer step. Must divide by `dp_size`. | | `--num-steps-per-rollout` | derived | Optimizer steps per rollout. | | `--micro-batch-size` | `1` | **Train pairs** per DiT forward (flat cut). | | `--micro-batch-size-sample` × `--micro-batch-size-tstep` | – | 2D tile per DiT forward. Set together; overrides `--micro-batch-size` with their product. | | `--num-rollout` | from dataset | Total rollout iterations. | | `--num-epoch` | – | Alternative to `--num-rollout`; ignored if both are set. | To check how these batch-related arguments interact and related to each other — see [the batch-knob invariant](/docs/diffusion/user-guide/concepts#the-batch-knob-invariant). ### Diffusion sampling | Flag | Default | What | | ------------------------------------------ | ------- | -------------------------------------------------------- | | `--diffusion-num-steps` | `10` | Denoising steps per rollout sample. | | `--diffusion-num-sde-steps` | `0` | How many steps become train pairs. `0` disables. | | `--diffusion-step-strategy-path` | – | Which steps. Overrides the bare count. | | `--diffusion-noise-level` | `0.7` | SDE noise injected during rollout. | | `--diffusion-guidance-scale` | `4.0` | CFG scale. | | `--diffusion-height` / `--diffusion-width` | `512` | Output resolution. | | `--diffusion-output-num-frames` | `1` | Frames — `1` for images. | | `--rollout-microgroup-size` | `1` | Samples per prompt sent in one `POST /rollout/generate`. | ### Precision | Flag | Default | What | | --------------------------- | ------- | ----------------------------------------------------------------------------------- | | `--diffusion-forward-dtype` | `bf16` | DiT forward compute — engine, FSDP `param_dtype`, and the training-side input cast. | | `--fsdp-master-dtype` | `fp32` | Sharded master weights and optimizer state. | | `--fsdp-reduce-dtype` | `fp32` | Gradient reduce-scatter. | See [Dtype Control](/docs/diffusion/advanced/dtype-control). ### Algorithm | Flag | Default | What | | -------------------------- | ------------- | ------------------------------------------------------------- | | `--loss-type` | `policy_loss` | `policy_loss` (Flow-GRPO), `nft`, `sft_loss`, `custom_loss`. | | `--advantage-estimator` | `grpo` | Only `grpo` today. | | `--diffusion-clip-range` | `1e-4` | PPO ratio clip. | | `--diffusion-adv-clip-max` | `5.0` | Advantage clamp. | | `--diffusion-kl-beta` | `0.0` | Reference-KL coefficient; `> 0` turns on a reference forward. | | `--globalize-reward-std` | off | Batch-wide std instead of per-group. | | `--globalize-reward-mean` | off | Batch-wide mean instead of per-prompt. | *** ## Complete reference ### Cluster | Flag | Type | Default | Notes | | -------------------------------------------- | --------- | ------- | ----------------------------------------------------------------------------------------------------- | | `--actor-num-nodes` | int | `1` | | | `--actor-num-gpus-per-node` | int | `8` | | | `--rollout-num-gpus` | int | – | For train-only SFT, unset colocates encoders with training; set it to reserve dedicated encoder GPUs. | | `--rollout-num-gpus-per-engine` | int | `1` | Like sglang's `tp_size`. | | `--num-gpus-per-node` | int | `8` | | | `--colocate` | flag | off | Required: CUDA IPC weight sync currently supports only colocated trainer and rollout execution. | | `--offload` | flag | off | `--offload-train` + `--offload-rollout`. | | `--offload-train` / `--no-offload-train` | tri-state | – | Always on under `--colocate`. | | `--offload-rollout` / `--no-offload-rollout` | tri-state | – | Always on under `--colocate`. | | `--distributed-backend` | str | `nccl` | | | `--distributed-timeout-minutes` | int | `10` | | ### Training backend | Flag | Type | Default | Notes | | --------------------------- | ----- | ------- | --------------------------------------------------------------------------------------------------------------------------------- | | `--train-backend` | enum | `fsdp` | Only value. | | `--fsdp-master-dtype` | enum | `fp32` | `fp32` / `bf16` / `fp16`. Load, shard, and optimizer-state precision. | | `--fsdp-reduce-dtype` | enum | `fp32` | `bf16` matches flow\_grpo's all-bf16 policy but adds cross-rank add-noise. | | `--diffusion-forward-dtype` | enum | `bf16` | `bf16` / `fp16` / `fp32`. | | `--fsdp-cpu-offload` | flag | off | Offloads params, grads, optimizer state; the optimizer then runs on CPU. | | `--fsdp-cpu-backend` | str | `gloo` | CPU process group for the above. | | `--dp-replicate-size` | int | `1` | Hybrid sharding: replica count. `dp_shard` takes the ranks left over from this and SP. | | `--sequence-parallel-size` | int | `1` | USP = Ulysses × Ring. | | `--ulysses-degree` | int | `0` | `0` = auto (Ulysses fills SP). Ring degree > 1 needs torch ≥ 2.11 and a ring-capable attention backend. | | `--fsdp-attention-backend` | str | – | diffusers `set_attention_backend` value. | | `--fsdp-flow-shift` | float | – | Training-side sigma grid shift, regenerated when no engine supplies scheduler meta (SFT). Distinct from `--diffusion-flow-shift`. | | `--gradient-checkpointing` | flag | off | | | `--deterministic-mode` | flag | off | See [Deterministic Training](/docs/diffusion/advanced/deterministic). | | `--train-env-vars` | JSON | `{}` | Extra env for the training processes. | ### Optimizer and schedule | Flag | Type | Default | Notes | | ----------------------------------------------- | ----- | --------------- | ----------------------------------------------- | | `--optimizer` | enum | `adam` | AdamW. Only value. | | `--lr` | float | `1e-6` | | | `--adam-beta1` / `--adam-beta2` | float | `0.9` / `0.999` | β₂ matches flow\_grpo, not the LLM-side `0.95`. | | `--adam-eps` | float | `1e-8` | | | `--weight-decay` | float | `0.0` | | | `--clip-grad` | float | `1.0` | | | `--lr-decay-style` | enum | `constant` | | | `--min-lr` / `--lr-warmup-init` | float | `0.0` | | | `--lr-warmup-iters` | int | `0` | | | `--lr-warmup-fraction` | float | – | | | `--lr-decay-iters` | int | – | | | `--lr-wsd-decay-iters` / `--lr-wsd-decay-style` | – | – | Warmup-stable-decay. | | `--use-checkpoint-lr-scheduler` | flag | on | | | `--override-lr-scheduler` | flag | off | | | `--seed` | int | `1234` | | ### Rollout | Flag | Type | Default | Notes | | ----------------------------------------------- | --------- | ----------------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------- | | `--hf-checkpoint` | str | – | **Required.** Pipeline to train and to serve; also the family source. | | `--diffusion-model-family` | str | – | Registered family key: `sd3`, `wan2_2`, `ltx`, `qwen_image`, `cosmos3`. Overrides name matching. | | `--rollout-function-path` | str | `miles.rollout.sglang_rollout.generate_rollout` | Generic Miles default. Diffusion recipes explicitly set `miles.rollout.sglang_diffusion_rollout.generate_rollout`. | | `--train-pipeline-config-path` | str | – | Your own `TrainPipelineConfig` for an unregistered family. Mutually exclusive with `--diffusion-model-family`. | | `--model-backend-path` | str | – | Override the family's model loader. | | `--diffusion-num-steps` | int | `10` | | | `--diffusion-flow-shift` | float | – | Generation-schedule shift, sent to the engine. | | `--rollout-microgroup-size` | int | `1` | | | `--diffusion-fps` | float | – | Video only. | | `--diffusion-output-num-frames` | int | `1` | | | `--diffusion-guidance-scale` | float | `4.0` | | | `--diffusion-guidance-scale-2` | float | – | Wan2.2 low-noise expert; **required** when training it. | | `--diffusion-true-cfg-scale` | float | – | | | `--diffusion-negative-prompt` | str | – | Defaults to `" "` on the engine when CFG is on. | | `--diffusion-noise-level` | float | `0.7` | | | `--diffusion-height` / `--diffusion-width` | int | `512` | Rollout output size; SFT center-crop size. | | `--diffusion-sde-type` | enum | `sde` | `sde` / `cps` / `ode`. Selects the train-side SDE backend too. | | `--sde-step-backend-path` | str | – | Custom dynamics. See [SDE backends](#sde-step-backends). | | `--diffusion-num-sde-steps` | int | `0` | | | `--diffusion-sde-window-range` | `"lo,hi"` | – | For `sde_window`. Defaults to `[0, num_inference_steps)`. | | `--diffusion-sde-candidate-steps` | `"1,2,3"` | – | Required by `epoch_global_random_choice`. | | `--diffusion-step-strategy-path` | str | – | Overrides the bare `--diffusion-num-sde-steps` selection. | | `--diffusion-log-prob-no-const` | flag | off | Drop log-prob constants on the engine (pairs with the CPS backend). | | `--diffusion-generator-device` | str | `cuda` | | | `--rollout-patch-group` | str | – | Comma-separated numeric-parity patch groups, e.g. `sgld`, `ltx`. | | `--update-weight-target-module` | str | `transformer` | Modules to train and sync. Wan2.2: `transformer,transformer_2`. | | `--update-weight-buffer-size` | int | 512 MiB | Weight-sync chunk size in bytes. | | `--rollout-seed` | int | `42` | | | `--over-sampling-batch-size` | int | – | Must equal `--rollout-batch-size` today. | | `--sglang-server-concurrency` | int | `512` | Per-engine in-flight request cap. | | `--use-distributed-post` | flag | off | Rollout HTTP POSTs go through per-node Ray actors instead of the local client (`MILES_HTTP_POST_ACTORS_PER_NODE` sets the count). | | `--use-miles-router` | flag | off | **Required** — the SGLang router is not supported here. | | `--miles-router-timeout` | float | – | | | `--miles-router-max-connections` | int | – | | | `--miles-router-health-check-failure-threshold` | int | `3` | | ### Data and batching | Flag | Type | Default | Notes | | ---------------------------------- | ---- | -------------- | --------------------------------------------------------------------------------------------------- | | `--prompt-data` | str | – | jsonl, one row per prompt. | | `--input-key` | str | `input` | | | `--metadata-key` | str | `metadata` | | | `--data-source-path` | str | – | Defaults to `RolloutDataSourceWithBuffer`. | | `--disable-rollout-global-dataset` | flag | off | Manage data yourself. | | `--rollout-batch-size` | int | – | | | `--n-samples-per-prompt` | int | `1` | | | `--global-batch-size` | int | derived | | | `--num-steps-per-rollout` | int | derived | | | `--micro-batch-size` | int | `1` | Flat: train-pair dicts per DiT forward, contiguous within an optimizer window. | | `--micro-batch-size-sample` | int | – | 2D: samples per DiT-forward tile. | | `--micro-batch-size-tstep` | int | – | 2D: SDE timesteps per tile. Set with the above. | | `--diffusion-train-iter-order` | enum | `sample_major` | `sample_major` / `timestep_major`. Tile visit order; only meaningful with the 2D pair. | | `--train-dp-split-mode` | enum | `contiguous` | `contiguous` lets a micro-batch reproduce a rollout microgroup exactly; `stride` deals round-robin. | | `--num-rollout` | int | – | | | `--num-epoch` | int | – | | | `--start-rollout-id` | int | – | Resumed from `--load` when unset. | | `--sft-encoder-checkpoint` | str | – | SFT only: tokenizer/text\_encoder/vae source. | | `--sft-frame-stride` | int | `1` | SFT encode temporal stride. | ### Evaluation | Flag | Type | Default | Notes | | ----------------------------- | ---- | ------- | ---------------------------------------------------- | | `--eval-interval` | int | – | Requires configured eval datasets. | | `--eval-prompt-data` | str+ | – | Repeatable ` ` pairs. | | `--eval-config` | str | – | OmegaConf YAML/JSON; overrides `--eval-prompt-data`. | | `--eval-function-path` | str | – | Defaults to `--rollout-function-path`. | | `--eval-input-key` | str | – | | | `--n-samples-per-eval-prompt` | int | `1` | | | `--diffusion-eval-num-steps` | int | – | Defaults to `--diffusion-num-steps`. | | `--skip-eval-before-train` | flag | off | | | `--eval-reward-key` | str | – | Defaults to `--reward-key`. | ### Algorithm and loss | Flag | Type | Default | Notes | | -------------------------------------- | ----- | ------------- | ------------------------------------------------------------------------------------------------------- | | `--loss-type` | enum | `policy_loss` | `policy_loss` / `nft` / `sft_loss` / `custom_loss`. | | `--advantage-estimator` | enum | `grpo` | | | `--disable-grpo-std-normalization` | flag | off | Dr.GRPO. Also forced off when `n_samples_per_prompt == 1`. | | `--globalize-reward-mean` | flag | off | Leave **off** for flow\_grpo parity. | | `--globalize-reward-std` | flag | off | **On** for flow\_grpo's PickScore recipe. | | `--diffusion-clip-range` | float | `1e-4` | | | `--diffusion-adv-clip-max` | float | `5.0` | Under `nft` this also sets the advantage-to-`r` slope. | | `--diffusion-recompute-old-log-prob` | flag | off | Recompute old log-probs with the trainer forward instead of trusting the rollout's. `policy_loss` only. | | `--diffusion-kl-beta` | float | `0.0` | | | `--ref-mode` | enum | – | `none` / `lora_base` / `ema`. Auto: `lora_base` when KL > 0, `ema` under `nft`. | | `--custom-prepare-train-batch-path` | str | – | Builds DiT inputs. | | `--custom-loss-function-path` | str | – | Loss **formula** only — the DiT forward stays in the actor. | | `--diffusion-nft-beta` | float | `1.0` | | | `--diffusion-nft-timestep-fraction` | float | `0.99` | | | `--no-diffusion-nft-adaptive-weight` | flag | off | | | `--no-diffusion-nft-shuffle-timesteps` | flag | off | | ### Reward | Flag | Type | Default | Notes | | ----------------------------------- | ----- | ------- | ------------------------------------------------------------------------------------------------ | | `--rm-type` | enum | – | `pickscore` / `ocr`. Overridable per sample via `metadata.rm_type`. | | `--reward-key` | str | – | When the reward is a dict. | | `--group-rm` | flag | off | Score a whole prompt group at once. | | `--custom-rm-path` | str | – | `async def rm(args, samples) -> list[float]`. Batched only. | | `--custom-reward-post-process-path` | str | – | Replace advantage normalisation. | | `--colocate-reward` | flag | off | Reward actors onto rollout GPUs (train 0.7 + rollout 0.25 + reward 0.05). Requires `--colocate`. | | `--pickscore-model-path` | str | – | Required for `--rm-type pickscore`. | | `--pickscore-processor-path` | str | – | Required for `--rm-type pickscore`. | | `--pickscore-num-workers` | int | `1` | | | `--pickscore-num-gpus-per-worker` | float | `1.0` | Fractional values allowed. | | `--pickscore-batch-size` | int | `8` | | | `--pickscore-num-frames` | int | – | Frames scored per video; unset = all. | | `--ocr-num-workers` | int | `4` | | | `--rollout-parser-num-workers` | int | `1` | Ray actors deserializing rollout responses. Raise when trajectory tensors are large. | ### Rollout customization hooks Every one takes a dotted path. | Flag | Replaces | | --------------------------------------------- | --------------------------------------------------------------------------------------------------------- | | `--custom-generate-function-path` | The inner `generate` call in the diffusion rollout. | | `--dynamic-sampling-filter-path` | Per-group keep/drop (DAPO-style). | | `--rollout-sample-filter-path` | Per-sample loss mask; set `sample.remove_sample = True`. Samples still count for advantage normalisation. | | `--buffer-filter-path` | Buffer selection. | | `--custom-convert-samples-to-train-data-path` | The whole samples → train-data conversion. | | `--custom-expand-samples-to-train-pairs-path` | Just the sample → train-pair expansion. | | `--custom-rollout-log-function-path` | Train rollout logging. | | `--custom-eval-rollout-log-function-path` | Eval rollout logging. | ### LoRA | Flag | Type | Default | Notes | | ------------------------------ | ---- | ----------- | ------------------------------------------------------------------------------ | | `--use-lora` | flag | off | | | `--lora-rank` / `--lora-alpha` | int | `64` / `64` | | | `--lora-target-modules` | str+ | – | Defaults per model family. | | `--lora-init-weights` | str | `gaussian` | `kaiming-uniform` maps to PEFT's default; other PEFT schemes pass through. | | `--lora-ipc-weight-sync` | flag | off | Push only `lora_A`/`lora_B`; the engine merges locally. Requires `--use-lora`. | ### EMA | Flag | Type | Default | Notes | | ------------------------ | ----- | ------- | ------------------------------------------------------------------------------------------------- | | `--use-ema` | flag | off | Maintains an EMA copy as πₒₗd. Needs a consumer (`--ref-mode ema` or `--ema-rollout-policy ema`). | | `--ema-rollout-policy` | enum | `live` | `live` / `ema`: which weights get pushed to rollout. | | `--ema-decay-init` | float | `0.001` | Decay during the flat period. | | `--ema-decay-ramp` | float | `0.001` | Per-step increase after the flat period; the ramp restarts from zero. | | `--ema-decay-max` | float | `0.5` | Ceiling. | | `--ema-decay-flat-steps` | int | `0` | | ### Checkpointing | Flag | Type | Default | Notes | | ----------------------------------- | ---- | ------- | ------------------------------------------------ | | `--save` | str | – | | | `--save-interval` | int | – | Requires `--save`. | | `--no-save-optim` | flag | off | Smaller checkpoints, but no resumption. | | `--load` | str | – | | | `--ckpt-step` | int | – | Defaults to `latest_checkpointed_iteration.txt`. | | `--no-load-optim` / `--no-load-rng` | flag | off | | ### Logging | Flag | Type | Default | Notes | | ------------------------------------------------------------------------------------ | ---- | ------------------- | --------------------------------------------------------------------------- | | `--use-wandb` | flag | off | | | `--wandb-project` | str | – | | | `--wandb-group` / `--wandb-run-id` / `--wandb-team` / `--wandb-host` / `--wandb-key` | str | – | | | `--wandb-mode` | enum | – | `online` / `offline` / `disabled`. Overrides `WANDB_MODE`. | | `--wandb-dir` | str | – | Defaults to `./wandb`. | | `--disable-wandb-random-suffix` | flag | off | Run names include a random suffix by default; pass this flag to disable it. | | `--wandb-log-num-images` | int | `0` | Images/videos per rollout; `0` disables. | | `--wandb-log-image-interval` | int | `1` | Send media every N rollouts. | | `--use-miles-dashboard` | flag | off | Async phase/trajectory telemetry. | | `--miles-dashboard-workspace` | str | `./miles_dashboard` | | ### Fault tolerance | Flag | Type | Default | Notes | | ----------------------------------- | ----- | ------- | ------------------------------------------------------- | | `--use-fault-tolerance` | flag | off | Restart dead engines during rollout. | | `--rollout-health-check-interval` | float | `30.0` | Raise it for video models — a request can take minutes. | | `--rollout-health-check-timeout` | float | `30.0` | | | `--rollout-health-check-first-wait` | float | `0` | Grace period for compilation/init. | ### Debugging | Flag | Type | Default | Notes | | ------------------------------------- | ------------- | ------- | -------------------------------------------------------------------------------------------------- | | `--debug-rollout-only` | flag | off | Rollout, no training. | | `--train-only` | flag | off | No engines, no weight sync, no eval. (`--debug-train-only` is a legacy alias.) | | `--debug-skip-optimizer-step` | flag | off | No backward/step — weights never drift. Use it to measure pure forward divergence from the engine. | | `--save-debug-rollout-data` | path template | – | `.format(rollout_id)`. | | `--load-debug-rollout-data` | path template | – | Replays a rollout with **no engines started**. | | `--load-debug-rollout-data-subsample` | float | – | | | `--save-debug-train-data` | path template | – | | | `--dump-details` | dir | – | Sets the debug dump paths for post-hoc analysis. | | `--diffusion-debug-mode` | flag | off | Engine returns per-step debug tensors; enables the `model_output_*_abs_diff` metrics. | `--debug-skip-optimizer-step` + `--diffusion-debug-mode` is the standard train/rollout alignment probe: frozen weights, and the metrics report exactly how far the two forwards drift. ### CI `--ci-test`, `--ci-metric-checker-key`, `--ci-metric-checker-threshold`. Used by `tests/e2e`; not meant for manual runs. *** ## SDE step backends `--diffusion-sde-type` picks the train-side backend automatically: | `--diffusion-sde-type` | Backend | Dynamics | | ---------------------- | ------------------------- | --------------------------------------------------------------- | | `sde` | `DiffusersSdeStepBackend` | Flow-matching SDE over diffusers scheduler sigmas. | | `ode` | `DiffusersSdeStepBackend` | Same class; deterministic rollout (used by NFT). | | `cps` | `CpsSdeStepBackend` | CPS kernel, σ = timestep ÷ divisor, log-prob without constants. | `--sde-step-backend-path` overrides the mapping with your own `SdeStepBackend` subclass.