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.

· 2 min read
Slurm GPU Scheduling: A Complete Guide 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
#!/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

MetricSource
GPU utilizationdcgm_exporter → DCGM_FI_DEV_GPU_UTIL
GPU memorydcgm_exporter → DCGM_FI_DEV_FB_USED
GPU temperaturedcgm_exporter → DCGM_FI_DEV_GPU_TEMP
Per-job GPU usageslurm_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.