Multi-Node Slurm GPU Script Generator

Why Generic Slurm Scripts Fail at Multi-Node Scale

A boilerplate batch script copied from an introductory university cluster wiki works for a single GPU job, but it breaks down rapidly when scaling across 8 to 512 GPUs. Distributed training workloads fail in four repeatable ways when hardware topology and Slurm resource isolation are left at their defaults.

Preemption and Timeout Signal Handling with SIGUSR1

On high-utilization clusters, jobs are terminated abruptly either by reaching their --time wallclock limit or through preemption on spot and scavenging partitions. When Slurm kills a job with SIGTERM followed 30 seconds later by SIGKILL, distributed frameworks rarely have enough time to synchronize gradients and flush multi-gigabyte sharded checkpoints to Lustre or WEKA storage.

Adding #SBATCH --signal=B:USR1@180 instructs the Slurm controller to dispatch SIGUSR1 to the batch script 180 seconds before hard termination. Two implementation details are mandatory for this mechanism to work reliably:

• The B: Batch Flag: Without the B: prefix, Slurm signals active srun job steps directly or drops the signal if substeps are transitioning. Specifying B:USR1@180 guarantees delivery to the top-level bash wrapper shell.
• Backgrounding srun with wait: Bash does not execute trap handlers while blocked on a foreground child process. Your script must launch srun ... & in the background and block on wait $! so the shell immediately catches SIGUSR1 and forwards it to the Python training processes.

import os
import signal
import torch
import torch.distributed as dist

preemption_requested = False

def handle_sigusr1(signum, frame):
    global preemption_requested
    preemption_requested = True

signal.signal(signal.SIGUSR1, handle_sigusr1)

# Inside your training step loop:
if preemption_requested:
    # Synchronize decision across all ranks so no rank enters the next collective alone
    flag_tensor = torch.tensor([1], dtype=torch.int32, device="cuda")
else:
    flag_tensor = torch.tensor([0], dtype=torch.int32, device="cuda")

dist.all_reduce(flag_tensor, op=dist.ReduceOp.MAX)

if flag_tensor.item() == 1:
    save_sharded_checkpoint(model, optimizer, step)
    dist.barrier()
    dist.destroy_process_group()
    raise SystemExit(0)

Always broadcast or all_reduce the signal flag across the process group before exiting. Because OS signal delivery across dozens of compute nodes occurs asynchronously over several milliseconds, acting on SIGUSR1 locally without cross-rank consensus causes early ranks to exit while trailing ranks enter the next forward pass and hang in an orphaned collective.

Interconnect-Specific NCCL Configuration: InfiniBand NDR vs RoCE v2 vs AWS EFA

Collective communication libraries require fundamentally different transport flags depending on the physical fabric connecting your GPU nodes. Applying InfiniBand verbs flags to an AWS EFA cluster or running RoCE v2 without explicit GID and traffic class pinning degrades AllReduce and AllGather bandwidth from 400 Gbps down to TCP fallback speeds.

Related Cluster Sizing and Fault Diagnosis Resources

Pair your Slurm launch script with accurate memory sizing and automated telemetry triage before and during your training runs.