Slurm GPU Scheduling: A Complete Guide for AI Workloads
Complete guide to Slurm GPU scheduling: MIG, MPS, GRES configuration. PyTorch and TensorFlow distributed training with Slurm scripts. HPC cluster optimization for AI workloads.
AI and deep learning workloads have fundamentally different resource requirements than traditional HPC simulations. GPU memory, NVLink bandwidth, GPU-to-GPU communication topology, and data loading pipelines make GPU scheduling far more complex than CPU scheduling.
Slurm GPU Basics
Slurm manages GPUs through the GRES (Generic Resource Scheduling) plugin:
# slurm.conf
GresTypes=gpu
# gres.conf
NodeName=gpu01 Name=gpu Type=h100 File=/dev/nvidia[0-7] Count=8
# Request 4 H100 GPUs
#SBATCH --gres=gpu:h100:4
#SBATCH --partition=gpu
MIG (Multi-Instance GPU)
NVIDIA H100 and A100 GPUs support MIG, partitioning a single physical GPU into up to 7 independent instances:
nvidia-smi -i 0 -mig 1
nvidia-smi mig -cgi 9,9,9,9,9,9,9 -C
MIG can boost GPU utilization from 30% to 90% for small inference and light training workloads by running 7 jobs simultaneously on one GPU.
Distributed Training with PyTorch FSDP
#!/bin/bash
#SBATCH --partition=gpu
#SBATCH --nodes=4
#SBATCH --gres=gpu:h100:8
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=64
export NCCL_SOCKET_IFNAME=ib0
export MASTER_ADDR=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n1)
srun torchrun --nnodes=$SLURM_NNODES --nproc_per_node=8 train_fsdp.py
DeepSpeed ZeRO-3
#!/bin/bash
#SBATCH --nodes=2
#SBATCH --gres=gpu:h100:8
deepspeed --num_nodes=$SLURM_NNODES --num_gpus=8 train.py \
--deepspeed ds_config_zero3.json
Job Arrays for Hyperparameter Search
#!/bin/bash
#SBATCH --array=1-100
#SBATCH --gres=gpu:h100:1
HPARAMS="hparams_${SLURM_ARRAY_TASK_ID}.json"
python train.py --config ${HPARAMS}
Failure Recovery
#SBATCH --requeue # Auto-requeue on node failure
if [ -f "checkpoint_latest.pt" ]; then
python train.py --resume checkpoint_latest.pt
else
python train.py --from-scratch
fi
InfiniBand and NCCL Optimization
export NCCL_IB_DISABLE=0
export NCCL_SOCKET_IFNAME=ib0
export NCCL_IB_HCA=mlx5_0,mlx5_1
Misconfigured networking can slow training by 10×. Always benchmark with all_reduce_perf after setup.
GPU Monitoring Metrics
| Metric | Source |
|---|---|
| GPU utilization | dcgm_exporter → DCGM_FI_DEV_GPU_UTIL |
| GPU memory | dcgm_exporter → DCGM_FI_DEV_FB_USED |
| GPU temperature | dcgm_exporter → DCGM_FI_DEV_GPU_TEMP |
| Per-job GPU usage | slurm_exporter → slurm_job_gpu_utilization |
For enterprise GPU cluster deployment and Slurm optimization, see our GPU Cluster solution. For performance tuning, see our Slurm optimization guide.