GPU Job Examples

Note

These are just examples, and may not work without modification. If you have issues please submit an Helpdesk ticket.

PyTorch Training

#!/bin/bash
#SBATCH --job-name=pytorch-train
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=8
#SBATCH --gres=gpu:1
#SBATCH --mem=64G
#SBATCH --time=12:00:00
#SBATCH --partition=<valid_gpu_partition> 
module purge
module load cuda python
pip install torch torchvision

python train_pytorch.py

TensorFlow Training

#!/bin/bash
#SBATCH --job-name=tf-train
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=8
#SBATCH --gres=gpu:1
#SBATCH --mem=64G
#SBATCH --time=12:00:00
#SBATCH --partition=<valid_gpu_partition> 
module purge
module load cuda python
pip install tensorflow

python train_tensorflow.py

Multi-GPU Training

#!/bin/bash
#SBATCH --job-name=multi-gpu
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=16
#SBATCH --gres=gpu:4
#SBATCH --mem=128G
#SBATCH --time=24:00:00
#SBATCH --partition=<valid_gpu_partition> 
module purge
module load cuda python

# PyTorch DDP example
python -m torch.distributed.run --nproc_per_node=4 train.py

GPU Inference

#!/bin/bash
#SBATCH --job-name=gpu-inference
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=4
#SBATCH --gres=gpu:1
#SBATCH --mem=32G
#SBATCH --time=02:00:00
#SBATCH --partition=<valid_gpu_partition> 
module purge
module load cuda python

# Run inference
python run_inference.py --model large --device cuda

Interactive GPU Session

# Request an interactive GPU session
srun --pty --partition=<valid_gpu_partition> --gres=gpu:1 --time=04:00:00 /bin/bash

# Once in the session
nvidia-smi
jupyter notebook --port=8888