slurm-gpu-batch

name

slurm-gpu-batch

description

Writes and reviews Slurm sbatch scripts for GPU clusters including gres GPU requests, partitions, memory and CPU sizing, time limits, job arrays, logging, multi-GPU srun patterns, the discipline of only requesting a GPU when the workload genuinely uses one, and the preflight/sampling/audit pattern that verifies a GPU job actually used the GPU. Use when the user mentions sbatch, Slurm, HPC, GPU jobs, Unity or other campus clusters, preempt GPU partitions, hyperparameter sweeps, or batch submission of ML experiments.

Slurm GPU batch scripts

This skill follows the Agent Skills layout. Teachings below draw from Princeton Research Computing Slurm (including GPUs and job arrays), Unity HPC batch jobs, and this repo’s example sweep script.

When to use

Apply when drafting or fixing #SBATCH scripts for GPU work, sweeps, or cluster submission. Always reconcile flags with local cluster docs before the user submits.

CPU first, GPU only when needed

Most clusters audit GPU utilization and email accounts whose jobs request GPUs but do not use them. On Unity, a recent admin report on this account flagged that 75.9% of GPU jobs over 14 days never touched the GPU, and mean VRAM usage was 8.2% of the request. That wastes shared hardware, hurts queue priority, and earns admin warnings.

Decision flow before adding --gres=gpu:* to a script:

1. Does the code actually call into CUDA? Inference, plotting, table generation, small numerical scripts, and anything that sets JAX_PLATFORMS=cpu or runs torch.set_default_device("cpu") belongs on a CPU partition. Do not request a GPU.
2. Will the GPU sit at >30% utilization for most of the wall time? If not, the workload is host-bound and the GPU is wasted. Use CPU.
3. Right-size VRAM. Pick the smallest VRAM bucket (--constraint=vram16, vram24, ...) that fits measured peak VRAM. Asking for 80 GB to run a 1 GB model triples queue time for no reason.
4. Right-size host RAM. Default --mem should be ~2× the measured peak (MaxRSS from sacct), not a flat 32 GB. Habitually low-utilization --mem requests get throttled by the scheduler.

When uncertain, run a short profiling job on CPU first, measure with sacct and nvidia-smi, then size the real submission to fit.

Verifying that a GPU job actually used the GPU

A GPU script should make it impossible to silently fall back to CPU. Bake three checks into the script:

1. Preflight assertion. After environment setup, before the real work, abort if no GPU is visible. Adapt to the framework:

echo "host=$(hostname) CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-<unset>}"
nvidia-smi -L
# JAX:
python -c "import jax; d=jax.devices(); print('jax devices:', d); assert any(x.platform=='gpu' for x in d), 'JAX is NOT on GPU — abort'"
# PyTorch:
# python -c "import torch; assert torch.cuda.is_available() and torch.cuda.device_count()>0, 'CUDA not available — abort'"

This one line is what catches "the job asked for a GPU but the framework silently used CPU" — by far the most common cause of GPU-jobs-with-zero-utilization in admin audits.

2. Mid-run utilization sampler. Fork a background loop that logs nvidia-smi once per minute:

( while sleep 60; do nvidia-smi --query-gpu=utilization.gpu,utilization.memory,memory.used --format=csv,noheader; done ) &
SMI_PID=$!
trap 'kill $SMI_PID 2>/dev/null || true' EXIT

If every sample reads 0%, 0%, <100 MiB, the job is misbehaving regardless of what the framework reported.

3. Post-run sacct audit. After the job finishes, read GPU and memory utilization from accounting:

sacct -j <jobid> --format=JobID,State,Elapsed,MaxRSS,ReqMem,AllocTRES%40,TRESUsageInMax%60

Use this to drive the next submission:

• If gres/gpuutil in TRESUsageInMax is near zero, drop --gres and move to a CPU partition.
• If MaxRSS is below ~30% of ReqMem, lower --mem.
• If average GPU memory used is well below the VRAM bucket, drop to a smaller bucket.

Some sites also provide tools like Princeton's jobstats or Unity's web dashboard. Prefer those when available, but the sacct columns above are universally present.

Script shape

A batch script should do four things in order (see Princeton’s overview).

1. Resources
   #SBATCH lines for partition, GPUs, CPUs, memory, wall time, job name, stdout or stderr paths, optional mail, optional account or qos.

2. Environment
   module load, conda activate, cd to the project, export for threads or CUDA if needed.

3. Preflight (GPU jobs only)
   Echo CUDA_VISIBLE_DEVICES, run nvidia-smi -L, and assert from the framework that the GPU is visible (assert torch.cuda.is_available() or assert any(d.platform=='gpu' for d in jax.devices())). Abort the job if the check fails. See "Verifying that a GPU job actually used the GPU" below.

4. Work
   The actual command (python, uv run, srun, etc.).

GPU requests

• Request GPUs only if the program uses them. This is the single most important rule. Extra or unused GPUs waste quota, hurt priority, increase queue time without speedup, and on monitored clusters they generate admin warnings (Princeton GPU jobs). If unsure whether the code is GPU-bound, profile first on a CPU partition and inspect nvidia-smi/sacct from a short test job before committing to a GPU sweep.
• Request GPUs with #SBATCH --gres=gpu:N per node. For interactive debugging use salloc with the same idea, for example --gres=gpu:1 as in Princeton’s GPU interactive example.
• Do not request multiple GPUs unless the stack uses them (DDP, multi-process layout, etc.). When in doubt, one GPU and a scaling check later.
• Right-size VRAM. Pick the smallest --constraint=vramN bucket that fits measured peak VRAM. Asking for an 80 GB card to hold a 1 GB model triples queue time.
• Right-size host RAM. Default --mem should be ~2× measured peak RSS, not a flat large number. Audit MaxRSS with sacct after each new workload and shrink the request accordingly.
• Many GPU workloads need enough CPU cores for dataloading and framework overhead. Set --cpus-per-task or -c to match the site and the code (Princeton GPU jobs).
• To see how a GPU partition is used, sinfo -p <partition> is a standard first check (Princeton useful commands).

Time and memory

• Set --time to a realistic wall limit. Jobs are killed at the limit. Princeton recommends planning headroom on top of measured runtime (Princeton Slurm).
• Set memory explicitly (--mem, --mem-per-cpu, or site-specific --mem-per-gpu) when defaults are too small or unclear.
• Per Unity, options passed on the sbatch command line override the same options inside the script file.

Single-node ML default

For typical PyTorch or JAX training on one node:

• --nodes=1 and --ntasks=1 (or equivalent short forms your site documents).
• One line #SBATCH --gres=gpu:1 unless the user needs more GPUs.
• Enough CPUs for workers and BLAS or framework threads.

Multiple GPU processes in one job

If several GPU processes must share one allocation (not the same as one multi-GPU training job), Princeton shows matching --ntasks and --gres with one GPU per srun and wait at the end (Running multiple jobs in parallel). Prefer job arrays when tasks are the same script with different indices.

Job arrays

• Declare with #SBATCH --array=0-99 or ranges and lists as in Princeton job arrays.
• Throttle concurrency with %k, for example #SBATCH --array=0-199%16 (Princeton job arrays).
• Use unique stdout and stderr. For arrays use %A (array job id) and %a (task id) in -o and -e.
• mkdir -p the log directory at the start of the script body so the first task does not fail on a missing path. See local_disentanglement/experiments/run_identifiability_sweep.sh in this repo for %A_%a logs, set -euo pipefail, decoding SLURM_ARRAY_TASK_ID into hyperparameters, and grouping runs with SLURM_ARRAY_JOB_ID (for example wandb groups).
• Every task must write to distinct checkpoint and artifact paths. Otherwise tasks clobber each other.

Shell hygiene

• Use set -euo pipefail after SBATCH lines when bash allows it.
• Echo hostname, key env vars if useful (CUDA_VISIBLE_DEVICES), and resolved hyperparameters for reproducible logs.
• Use explicit cd to the repo root or absolute paths so the job does not depend on the submission directory by accident.

Site-specific checklist (confirm with user or cluster docs)

Before finalizing a script, align with the local site:

• Partition name (-p or --partition) and any QOS or account (-A, --qos).
• How GPU types are selected (--constraint, or gres with a type if supported).
• Wall time and memory limits per partition.
• Whether short flags (-c, -t, -J) or long forms are customary.
• How to load Python (modules, conda, uv run, container).
• Where to write large outputs (scratch versus home) and quota rules.

This repo’s run_identifiability_sweep.sh uses -p gpu-preempt, --gres=gpu:1, -c, --mem, and a constraint string. Treat that as one cluster’s pattern, not universal Slurm syntax.

Operations quick reference

Action	Command
Submit	sbatch script.sh
My jobs	squeue --me (Unity, Princeton)
Job details after run	sacct -j JOBID (Unity)
Cancel	scancel JOBID
Mail	#SBATCH --mail-type=BEGIN,END,FAIL and --mail-user=... (Unity)

On Princeton systems, jobstats gives post-run CPU, memory, and GPU metrics (Princeton useful commands). Use the equivalent your site documents if any.

Copy-paste templates

For minimal starting points (single GPU job, GPU array, two-GPU srun pattern), see references/REFERENCE.md.