PyTorch¶

PyTorch is an open-source deep learning framework providing tensor computation with GPU acceleration and automatic differentiation for building and training neural networks.

This page covers how to set up a PyTorch environment on the cluster and run training jobs at different scales: single GPU, multi-GPU on one node, and multi-node distributed training.

Setting up a PyTorch environment¶

Create a dedicated conda environment using the Miniforge3 module:

module load Miniforge3
conda deactivate

mamba create --prefix /scratch/$USER/conda-envs/pytorch python=3.12 -y
conda activate /scratch/$USER/conda-envs/pytorch

# Install PyTorch with CUDA support from the pytorch channel
mamba install pytorch torchvision torchaudio pytorch-cuda=12.4 -c pytorch -c nvidia

# Install any additional packages you need
mamba install transformers lightning wandb  # examples

Verify the installation

Request a quick interactive session with a GPU to confirm everything works:

salloc --qos=test --gres=gpu:1 --mem=8G --time=00:10:00 srun --pty bash
module load Miniforge3
conda activate /scratch/$USER/conda-envs/pytorch
python -c "import torch; print(torch.cuda.is_available(), torch.cuda.device_count(), 'GPUs')"

Single GPU¶

The simplest case: one GPU, one process:

pytorch.slurm

#!/bin/bash
#SBATCH --qos=regular
#SBATCH --job-name=pytorch
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=8
#SBATCH --gres=gpu:1
#SBATCH --constraint=rtx3090
#SBATCH --mem=32G
#SBATCH --time=24:00:00
#SBATCH --output=%x.%j.out
#SBATCH --error=%x.%j.err

module load Miniforge3
conda activate /scratch/$USER/conda-envs/pytorch

python train.py

In your training script, move the model and data to GPU:

device = torch.device("cuda")
model = MyModel().to(device)

Single node, multiple GPUs¶

Use PyTorch's DistributedDataParallel (DDP) with torchrun to train across multiple GPUs on the same node. This is the recommended approach over DataParallel.

pytorch_ddp.slurm

#!/bin/bash
#SBATCH --qos=regular
#SBATCH --job-name=pytorch_ddp
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=32
#SBATCH --gres=gpu:4
#SBATCH --constraint=a100-sxm4
#SBATCH --mem=200G
#SBATCH --time=24:00:00
#SBATCH --output=%x.%j.out
#SBATCH --error=%x.%j.err

module load Miniforge3
conda activate /scratch/$USER/conda-envs/pytorch

torchrun \
    --standalone \
    --nproc_per_node=${SLURM_GPUS_ON_NODE:-4} \
    train_ddp.py

A100 SXM4 nodes

The A100 SXM4 nodes have 8 GPUs connected via NVLink (600 GB/s), ideal for multi-GPU training. Jobs on the general partition require a minimum of 4 GPUs on these nodes. See GPU accelerators for details.

The training script needs minimal changes to support DDP:

Minimal DDP training script

import os
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler

def main():
    # torchrun sets these automatically
    dist.init_process_group("nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)

    # Model
    model = MyModel().to(local_rank)
    model = DDP(model, device_ids=[local_rank])

    # Data: DistributedSampler splits the dataset across GPUs
    dataset = MyDataset()
    sampler = DistributedSampler(dataset)
    loader = DataLoader(dataset, batch_size=64, sampler=sampler)

    # Training loop
    for epoch in range(num_epochs):
        sampler.set_epoch(epoch)  # shuffle differently each epoch
        for batch in loader:
            inputs, targets = batch
            inputs = inputs.to(local_rank)
            targets = targets.to(local_rank)

            outputs = model(inputs)
            loss = criterion(outputs, targets)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

    dist.destroy_process_group()

if __name__ == "__main__":
    main()

Multi-node, multiple GPUs¶

Scale training across multiple nodes using torchrun with SLURM environment variables. NCCL handles inter-node communication over InfiniBand automatically.

pytorch_multinode.slurm

#!/bin/bash
#SBATCH --qos=regular
#SBATCH --job-name=pytorch_multinode
#SBATCH --nodes=2
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=32
#SBATCH --gres=gpu:8
#SBATCH --constraint=a100-sxm4
#SBATCH --mem=400G
#SBATCH --time=24:00:00
#SBATCH --output=%x.%j.out
#SBATCH --error=%x.%j.err

module load Miniforge3
conda activate /scratch/$USER/conda-envs/pytorch

# Let NCCL use InfiniBand
export NCCL_IB_DISABLE=0
export NCCL_DEBUG=INFO

# Rendezvous via the first node
MASTER_ADDR=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n 1)
MASTER_PORT=29500

srun torchrun \
    --nnodes=${SLURM_NNODES} \
    --nproc_per_node=${SLURM_GPUS_ON_NODE:-8} \
    --rdzv_id=${SLURM_JOB_ID} \
    --rdzv_backend=c10d \
    --rdzv_endpoint=${MASTER_ADDR}:${MASTER_PORT} \
    train_ddp.py

The same train_ddp.py script from the single-node example works without changes; torchrun sets the rank and world size automatically.

Verify scaling before using multiple nodes

Multi-node training adds communication overhead. Always benchmark single-node first and only scale out when you have confirmed that training is GPU-bound and benefits from more devices.

Checkpointing

Save checkpoints regularly, especially on preemption-gpu where jobs can be killed and requeued. Save only on rank 0 to avoid file conflicts:

if dist.get_rank() == 0:
    torch.save(model.module.state_dict(), "checkpoint.pt")

PyTorch¶

Setting up a PyTorch environment¶

Single GPU¶

Single node, multiple GPUs¶

Multi-node, multiple GPUs¶

See also¶