Running on Slurm¶
Submit a distributed training job through TraceML using the existing Slurm
template. Slurm starts one TraceML launcher per node; each launcher starts
its local workers through torchrun.
Your script must already support distributed training and use a supported automatic trainer or explicit integration.
1. Prepare the job¶
The template contains two files in
examples/distributed/slurm/:
traceml_ddp.sbatch: requests resources and starts one task per node.launch.sh: expands each node's Slurm variables and launches training.
Before submitting:
- Adjust the node count, GPUs, CPUs, time limit, and any account or partition settings in the sbatch file for your cluster.
- Activate your training environment before
srun. TraceML and your framework must be available on every node. - Keep
--ntasks-per-node=1. This template uses torchrun to create the GPU workers; Slurm must not create another task for every GPU. - Put the scripts and log directory on shared storage. The template uses
./logsunder the submit directory; change--logs-dirinlaunch.shif needed.
The supplied wrapper runs examples/distributed/ddp_minimal.py, which already
contains explicit TraceML setup. Replace that path with your training script
and match its distributed configuration to the requested resources.
2. Submit¶
From the TraceML repository root:
sbatch examples/distributed/slurm/traceml_ddp.sbatch
The job uses two nodes with four GPUs each by default. Node 0 is the first host in the allocation and owns both torchrun rendezvous and the TraceML aggregator. Summary mode is the default for multi-node runs.
3. Read the result¶
After telemetry finalizes successfully, the report is saved under:
logs/ddp-<job-id>-<restart-count>/final_summary.json
logs/ddp-<job-id>-<restart-count>/final_summary.txt
Each node's training streams are saved under nodes/node_<node-rank>/ in that
run folder. Slurm also saves the batch output as traceml-<job-id>.out.
Add --html-report to the wrapper for an HTML report. See
Reading the Output, Compare Runs, or
Training Crashes for the next step.
How the template launches training¶
The sbatch file exports settings shared by every node:
export MASTER_ADDR="$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)"
export RUN_NAME="ddp-${SLURM_JOB_ID}-${SLURM_RESTART_COUNT:-0}"
srun examples/distributed/slurm/launch.sh
The wrapper expands the per-node settings where it runs:
exec traceml run examples/distributed/ddp_minimal.py \
--mode=summary \
--run-name="${RUN_NAME}" \
--nnodes="${SLURM_NNODES}" \
--node-rank="${SLURM_NODEID}" \
--nproc-per-node="${SLURM_GPUS_ON_NODE}" \
--master-addr="${MASTER_ADDR}" \
--master-port=29500
| TraceML setting | Slurm source |
|---|---|
--nnodes |
SLURM_NNODES |
--node-rank |
SLURM_NODEID, expanded on each node |
--nproc-per-node |
SLURM_GPUS_ON_NODE, the allocated GPU count |
--master-addr |
First host in SLURM_JOB_NODELIST |
--run-name |
Job ID plus restart count, shared by all nodes |
Use a fresh run name for each launch. The restart count distinguishes requeued jobs, which retain their job ID. TraceML refuses to overwrite an existing run folder.
Expand the node rank on each node¶
Keep the launch command in the wrapper. This inline command is incorrect:
# The batch shell expands the rank once, before srun launches the other nodes.
srun traceml run train.py --node-rank=$SLURM_NODEID ...
It can give every node rank 0. Running launch.sh once per node lets each
process read its own SLURM_NODEID.
Cluster adjustments¶
GPU allocation and environment¶
The template uses --gres=gpu:4; some clusters require
--gpus-per-node=4. Ensure each launcher sees all GPUs allocated to its node
and that SLURM_GPUS_ON_NODE provides the intended numeric worker count.
SLURM_GPUS_PER_NODE may include a GPU type and is not used by the wrapper.
For CPU training, set an explicit local worker count in the wrapper instead of deriving it from GPUs. Keep one launcher task per node.
If your site resets task environments, also activate the environment inside
launch.sh. Keep slow setup before the launch so node 0 can start the
aggregator promptly.
Network access¶
Other nodes must reach node 0 on two ports:
| Purpose | Default port |
|---|---|
| torchrun rendezvous | 29500 |
| TraceML telemetry | 29765 |
If the first host's name is not reachable, set MASTER_ADDR to a reachable
node-0 IP. For separate telemetry interfaces, custom ports, or bind addresses,
use the advanced distributed settings.
For NCCL connection problems, your cluster may require
NCCL_SOCKET_IFNAME=<interface>. NCCL_DEBUG=INFO helps inspect that setup.
Slow finalization¶
The finalization budget is 300 seconds by default. For slow shared storage or
late telemetry, add --finalize-timeout-sec <seconds> to the wrapper or set
TRACEML_FINALIZE_TIMEOUT_SEC. This controls telemetry shutdown, not Slurm's
job time limit.