miles/docs/diffusion/user-guide/launch-script.md

title, description

title	description
Launch Scripts	What a miles-diffusion launch script does when you run it, how it is structured, and the ways to override a recipe.

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:

1. 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"
   

2. Set a script option, as a flag or a MILES_SCRIPT_* env var.

3. 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

1. Kills stale sglang / ray / miles processes.
2. 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.
3. Builds the Ray runtime env: NCCL socket vars, MASTER_ADDR, PYTHONPATH, plus anything from --extra-env-vars.
4. 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):

32 prompts × 16 samples             = 512 samples per rollout
512 samples × 2 SDE steps           = 1024 train pairs
1024 pairs ÷ 4 DP ranks             = 256 pairs per rank
256 ÷ 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

Next

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