Move a torchrun sbatch script to launch¶

Replace the script's requeue guard, rendezvous id, checkpoint search and srun line with a launch.yaml, and make the job's command launch train. A job file for go/dispatch-slurm then needs nothing else from the script.

Start from what the script does¶

test/reference/train.sbatch is the kind of script this replaces. It's written from torchrun's own multi-node examples, and has four parts:

#SBATCH --requeue
trap requeue_once TERM
RDZV_ID="$SLURM_JOB_ID.$RESTARTS"
LATEST="$(ls -d "$OUTPUT_DIR"/checkpoint-* 2>/dev/null | sort -V | tail -n1 || true)"
srun --nodes="$SLURM_JOB_NUM_NODES" --ntasks-per-node=1 torchrun --rdzv-id="$RDZV_ID" ...

Write the launch.yaml¶

version: 1
train:
  command: [python, train.py, --output_dir, /shared/run]
  procsPerNode: 4
checkpoint:
  root: /shared/run
  framework: hf
  resumeFlag: --resume_from_checkpoint

train.command is your payload, without torchrun or any of its flags. Check the file with launch validate, and see what a run would do with launch train -print.

Set the job's command¶

Make the command in the job file [launch, train]. go/dispatch-slurm generates the batch script, which ends by running it.

The script's part	What takes its place
`--requeue`, the trap and the restart limit	The adapter submits with `--no-requeue`, and the supervisor decides on recovery
A rendezvous id from the restart count	`--rdzv-id` is `DISPATCH_ATTEMPT`, which is new for every attempt
The search for the highest `checkpoint-<step>`	One `publish resolve`, so every node resumes from the latest save that was published complete
`srun` and the first node's address	The launcher's coordinator starts an agent on each node with the same torchrun command

The launcher also reports why the run ended, which the script can't do. See the command line.

Limits¶

A checkpointing job on Slurm is refused for now. The grace depends on the site's process tracker, which go/dispatch-slurm doesn't export yet (DISPATCH_SLURM_PROCTRACK, spec 0001 D14). Until it does, this example stops with exit 64. Add signal: {grace: 0s} to run it anyway, relying on the payload's periodic saves, with no save on a signal.
gpu.health: required is checked on the batch step's node too. Under Slurm the coordinator opens NVML before the run, so a batch step that can't see a GPU when srun steps can is refused with exit 64. Under optional, the coordinator's document says absent while the agents read.
Neither the script nor this page's configuration has run on a real cluster yet. The script has only been tested with fake Slurm commands.