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 <file.yaml> 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 |
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. |
--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. |
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. |
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. |
--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. |
--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 <name> <path> 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.