Launch Scripts

Every recipe ships as a Python launch script under scripts/, and starting a run is one command:

python3 scripts/run_diffusion_grpo_pickscore_5gpu_flowgrpo_aligned.py

This page explains what that command does and how to change what it runs. For the meaning of individual flags, see the CLI Reference.

How a launch script starts a training job

A launch script is a recipe, not the training process. It assembles the full train_diffusion.py command line, starts a local Ray cluster, and submits the command as a Ray job — submitted rather than run directly, so the driver lives inside the cluster and sees every node’s GPUs.

Layer	Location	Role
Launch script	`scripts/run_*.py`	Holds the recipe: the flag blocks and tuned values
Command utilities	`miles/utils/external_utils/command_utils.py`	Starts Ray and submits the job
Training entrypoint	`train_diffusion.py`	The train loop (~90 lines), run inside the Ray job

The structure of a launch script

Every launcher follows the same layout:

@dataclass
class ScriptArgs(U.ExecuteTrainConfig):     # inherits output_dir, num_nodes, ...
    num_rollout: int = 400
    extra_args: str = ""                     # appended verbatim to the command line

def prepare(args) -> str:                    # download the dataset, return its local path
    ...

def execute(args, data_dir) -> None:
    ckpt_args      = "..."                   # named groups of train_diffusion.py flags
    rollout_args   = "..."
    ...
    U.execute_train(train_args=" ".join(groups), num_gpus_per_node=5, config=args)

@U.dataclass_cli
def main(args): execute(args, prepare(args))

ScriptArgs — script options as flags and `MILES_SCRIPT_*` env vars

@U.dataclass_cli exposes each ScriptArgs field twice: as a --kebab-case CLI option (--num-rollout) and as an environment variable with the MILES_SCRIPT_ prefix (MILES_SCRIPT_NUM_ROLLOUT). A command-line value beats the env var, which beats the field default. The env form is how wrappers and cluster tooling inject machine-specific values without editing the script. Options shared by every launcher, from the ExecuteTrainConfig base class and repo convention:

Option	Default	Purpose
`--num-nodes`	`$SLURM_JOB_NUM_NODES` or `1`	Training nodes
`--cuda-visible-devices`	inherited from the environment	Physical GPUs the job may use. Applied by exporting it for `ray start`.
`--output-dir`	`<repo>/logs`	Where checkpoints and dumps are written
`--extra-env-vars`	empty	Extra env vars added to the Ray runtime env
`--extra-args`	empty	Extra flags appended to the `train_diffusion.py` command line

wandb flags are emitted only when WANDB_API_KEY is set — logging turns on by exporting the key, with no script change.

execute() — the flag groups

The command line is assembled as one string block per concern, then concatenated — the parser sees a flat list, and the groups follow concern, not flag prefix (a --diffusion-* flag can live in rollout_args, a --micro-batch-size-* flag in perf_args).

Block	What it carries
`ckpt_args`	`--hf-checkpoint` (required — also selects the model family), `--save` / `--load`
`rollout_args`	Prompt data, batch shape, sampler: steps, guidance, noise level, SDE step strategy
`eval_args`	Eval datasets and cadence; eval requests are pure generation, no trajectory
`grpo_args`	Advantage estimator, reward normalisation, clip range
`optimizer_args`	LR, betas, weight decay
`lora_args`	Rank, targets, `--lora-ipc-weight-sync`
`reward_args`	`--rm-type` and the reward worker pool
`sglang_args`	Engine side: router, concurrency, `--sglang-*` passthrough
`train_backend_args`	The three dtype flags, `--update-weight-target-module`
`perf_args`	Gradient checkpointing, micro-batch tiling, parser workers
`misc_args`	GPU layout, `--colocate`, `--deterministic-mode`

Ways to override a recipe

From lightest to heaviest:

Append flags with --extra-args. The value lands at the end of the command line, and for an argparse flag given twice the later occurrence wins — so it overrides anything the recipe sets:
python3 scripts/run_diffusion_grpo_pickscore_5gpu_flowgrpo_aligned.py \ --num-rollout 50 --extra-args "--diffusion-kl-beta 0.02"
Set a script option, as a flag or a MILES_SCRIPT_* env var.
Edit the script. The launcher is the canonical home of a recipe’s hyperparameters; change the flag blocks directly for anything you want to keep.

The Typer flags are also the CI interface: an e2e test names the recipe and passes its knobs (script="scripts/run_diffusion_grpo_sd3_ocr_sglang.py", args=["--num-rollout", "2", ...]), so launch scripts must stay runnable with no arguments and configurable through ScriptArgs — CI drives them the same way you do.

What execute_train runs on your machine

Kills stale sglang / ray / miles processes.
Starts a fresh cluster with export CUDA_VISIBLE_DEVICES=... && ray start --head — the device list must be in the raylet’s own environment; set per job or per actor it never reaches the scheduler, which then places work on excluded GPUs.
Builds the Ray runtime env: NCCL socket vars, MASTER_ADDR, PYTHONPATH, plus anything from --extra-env-vars.
Submits the job: ray job submit -- python3 train_diffusion.py <flags>.

Env var	Effect
`MILES_SCRIPT_EXTERNAL_RAY=1`	A scheduler already built the Ray cluster: skip the teardown and `ray start`, only submit. `--cuda-visible-devices` must be empty here — export it before each `ray start` instead. Used by the multi-node recipe below.
`MILES_SCRIPT_ENABLE_RAY_SUBMIT=0`	Run everything except the submission — shows what a launcher would do
`MASTER_ADDR`	Where the single-node `ray start --head` binds, default `127.0.0.1`. The training process group does NOT use it: rank 0 negotiates its real node IP and a free port at runtime, so multi-node needs no setting here

The batch-size arithmetic

Diffusion adds a dimension LLM RL does not have: one sample expands into several train pairs, one per trained denoising step.

samples per rollout   = rollout_batch_size × n_samples_per_prompt
train pairs           = samples × (number of SDE step indices)

The trajectory-level knobs are locked by the batch-knob invariant; contradictory values abort at parse time. global_batch_size counts samples and must divide by dp_size (= train world size ÷ --sequence-parallel-size).

Knob	Unit	Meaning
`--num-steps-per-rollout`	optimizer steps	Windows the pairs are cut into per rollout
`--micro-batch-size`	train pairs	Flat: pairs per DiT forward
`--micro-batch-size-sample` × `--micro-batch-size-tstep`	samples × timesteps	2D tile per DiT forward; overrides the flat knob with `S × T`
`--diffusion-train-iter-order`	—	Tile visit order (2D only)
`--train-dp-split-mode`	—	`contiguous` or `stride` dealing of pairs to DP ranks

Most recipes use the 2D tile — it states “S samples × T timesteps per forward” directly, which the flat knob cannot express. Wan2.2 is the exception: --micro-batch-size 2 is its way of keeping each forward on one side of the high/low-noise expert boundary. Worked example (the Qwen-Image recipe):

prompts × 16 samples             = 512 samples per rollout
samples × 2 SDE steps           = 1024 train pairs
pairs ÷ 4 DP ranks             = 256 pairs per rank
÷ num_steps_per_rollout(2)      = 128 pairs per optimizer step per rank
tile = 8 samples × 1 timestep       = 8 pairs per DiT forward → 16 forwards per step

Two failure modes to recognize:

Error	Fix
`num_pairs_shard=N not divisible by num_steps_per_rollout=M`	Adjust `rollout_batch_size`, `n_samples_per_prompt`, or `num_steps_per_rollout` so the per-rank pair count splits evenly
`Micro-batch mixes denoising phases` (Wan2.2)	A micro-batch spans the expert boundary; set `--micro-batch-size 1` or a phase-pure tiling

Adapting a recipe

Start from the closest existing script and change one group at a time:

Goal	What to change
Different reward	`reward_args` — swap `--rm-type`, or point `--custom-rm-path` at your own `async def rm(args, samples) -> list[float]`
Different prompts	`--prompt-data`, `--input-key`, and the `prepare()` download
Different checkpoint	`--hf-checkpoint`, plus `--diffusion-model-family` if the name carries no family hint
Smaller GPU budget	Lower `--rollout-batch-size`, `--n-samples-per-prompt`, `--rollout-microgroup-size`; keep the batch identity satisfied
Train more denoising steps	Raise `--diffusion-num-sde-steps` — this multiplies train pairs
Try a flag without editing the file	`--extra-args "--flag value"`

Multi-node training

The worked example is scripts/run_diffusion_grpo_wan22_pickscore_17gpu_multinode.py: wan2.2-A14B full finetune on 2 nodes × 8 GPUs (train + rollout colocated) plus 1 reward GPU on a separate node. Multi-node runs submit into a cluster you build yourself; the launcher only submits (MILES_SCRIPT_EXTERNAL_RAY=1).

Bring up the cluster

Each node runs ONE of the blocks below, chosen by its role. The two export lines must be in the environment of the ray start daemon itself, which is why every block repeats them. On the head node:

ulimit -n 1000000
export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
ray start --head --port=6379 --num-gpus 8 --dashboard-host 0.0.0.0

On every other training node:

ulimit -n 1000000
export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
ray start --address=<head-ip>:6379 --num-gpus 8

On the reward node:

ulimit -n 1000000
export CUDA_VISIBLE_DEVICES=0
ray start --address=<head-ip>:6379 --num-gpus 1

ulimit -n: a non-interactive ssh shell defaults to 1024 open files. The raylet inherits it and, once the FSDP actors and engines connect, dies with epoll: Too many open files — but the driver log shows a misleading ActorUnavailableError: ... RpcError: Socket closed. Verify with grep "Max open files" /proc/$(pgrep -f raylet | head -1)/limits after ray start.
CUDA_VISIBLE_DEVICES follows the same rule the launcher automates single-node: it must be in the raylet’s environment, and with an external cluster the launcher refuses a --cuda-visible-devices of its own.

ray status on the head should show every node and the full GPU count.

Submit

ulimit -n 1000000   # the driver needs it too
MILES_SCRIPT_EXTERNAL_RAY=1 python3 scripts/run_diffusion_grpo_wan22_pickscore_17gpu_multinode.py

Reward workers (--pickscore-num-workers 4 --pickscore-num-gpus-per-worker 0.25, no --colocate-reward) are default-scheduled and land on the only free GPU: the reward node.

Verify the run is healthy

With --rollout-patch-group wan and --sglang-attention-backend torch_sdpa, train/model_output_mean_abs_diff is expected to be exactly 0.0 from the first optimizer step. The 4-GPU proxy E2E standard records 0.0, and the documented 17-GPU runs sustained it over 200 rollouts. Any nonzero value indicates a train/rollout parity regression; check the patch group, backend, dtype, and versions on every rank.
Both training nodes near 100% GPU util during rollout; a 100%-vs-idle split between nodes means the reward workers were packed onto one node.

Multi-node pitfalls

Pitfall	What to know
Stale engine hijacks the port	Killing a driver by name leaves its spawned engine alive; the next driver’s health check talks to the old, unpatched server while the new one dies on `EADDRINUSE`. Between runs kill by GPU pid and check `nvidia-smi` shows no compute processes
Version strings lie under editable installs	`pip show sglang` reports the install-time commit; trust `git rev-parse HEAD` in the checkout (or the `-e git+...@<sha>` line in `pip freeze`)
Weight-sync bucket	`--update-weight-buffer-size 4294967296` (4 GiB). In one documented A14B run, 512 MiB took ~92 s and 4 GiB ~15 s; these are run-specific observations, not a benchmark guarantee.
Determinism env vars are trainer-side	`--deterministic-mode` sets `NCCL_DETERMINISTIC` and `CUBLAS_WORKSPACE_CONFIG` on the FSDP actors; check rollout parity with `train/model_output_mean_abs_diff`

CLI Reference — every flag, grouped.
Core Concepts — what one rollout iteration does, object by object.
Dtype Control — what the three dtype flags actually do.

Getting Started

User Guide

Models

Advanced

Developer

Launch Scripts

How a launch script starts a training job

The structure of a launch script

ScriptArgs — script options as flags and `MILES_SCRIPT_*` env vars

execute() — the flag groups

Ways to override a recipe

What execute_train runs on your machine

The batch-size arithmetic

Adapting a recipe

Multi-node training

Bring up the cluster

Submit

Verify the run is healthy

Multi-node pitfalls

Next

​How a launch script starts a training job

​The structure of a launch script

​ScriptArgs — script options as flags and MILES_SCRIPT_* env vars

​execute() — the flag groups

​Ways to override a recipe

​What execute_train runs on your machine

​The batch-size arithmetic

​Adapting a recipe

​Multi-node training

​Bring up the cluster

​Submit

​Verify the run is healthy

​Multi-node pitfalls

​Next

How a launch script starts a training job

The structure of a launch script

ScriptArgs — script options as flags and `MILES_SCRIPT_*` env vars

execute() — the flag groups

Ways to override a recipe

What execute_train runs on your machine

The batch-size arithmetic

Adapting a recipe

Multi-node training

Bring up the cluster

Submit

Verify the run is healthy

Multi-node pitfalls

Next