miles/docs/diffusion/user-guide/cli-reference.md

title, description

title	description
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 <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

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.

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.

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.