PyTorch on DeltaAI

DeltaAI provides pre-built PyTorch environments as modules and NVIDIA NGC containers. The recommended approach is to load a module, which gives you a tested conda environment with GPU support, Slingshot 11 networking, and common data-science packages ready to go.

Pre-Built Modules (Recommended)

Available versions

Module	PyTorch	CUDA	Python	Notes
python/miniforge3_pytorch/2.12.0	2.12.0	13.0	3.12	Newest – CUDA 13 toolchain
python/miniforge3_pytorch/2.11.0	2.11.0	13.0	3.12	First CUDA 13 build
python/miniforge3_pytorch/2.10.0	2.10.0	12.9	3.12	Default – stable release, production-tested
python/miniforge3_pytorch/2.7.0	2.7.0	12.6	3.12	Development build
python/miniforge3_pytorch/2.5.0	2.5.0	12.4	3.10	Legacy

The CUDA column is the version bundled in the PyTorch wheel. The 2.11.0 and 2.12.0 modules additionally load the site cuda/13.1.1 module for CUDA_HOME and nvcc; see Building Custom CUDA Extensions.

module load python/miniforge3_pytorch            # loads the default (2.10.0)
conda activate base                              # activate the conda environment

The unversioned name resolves to the current default. Name a version explicitly – module load python/miniforge3_pytorch/2.12.0 – to pin your job against a future default change.

Each module includes jupyter, jupyterlab, matplotlib, pandas, scipy, scikit-learn, mpi4py, tensorboard, wandb, and nvitop.

Note

Always run conda activate base after loading the module. The module puts the environment on your PATH, but conda activate runs additional setup scripts that configure library paths needed by some packages.

Tip

For Jupyter notebook support through Open OnDemand, run setup-kernel after loading the module to register a Jupyter kernel.

Quick verification (run on a compute node):

module load python/miniforge3_pytorch/2.10.0
conda activate base
python -c "import torch; print(torch.__version__, torch.cuda.is_available())"
# Expected: 2.10.0+cu129 True

Single-GPU Job

#!/bin/bash
#SBATCH --account=account_name
#SBATCH --partition=ghx4
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=16
#SBATCH --gpus-per-node=1
#SBATCH --time=00:20:00
#SBATCH --job-name=pytorch-single

# Environment: set explicitly so the job does not depend on your submit shell
module load python/miniforge3_pytorch/2.12.0
conda activate base

srun python3 my_script.py

Multi-Node Distributed Training

Use torchrun with the NCCL backend for DistributedDataParallel (DDP) training across multiple nodes. The nccl-ofi-plugin module is part of the default login environment, so the AWS OFI NCCL plugin that bridges NCCL to Slingshot 11 is already configured – no additional module loads and no NCCL environment variables are needed.

#!/bin/bash
#SBATCH --account=account_name
#SBATCH --partition=ghx4
#SBATCH --nodes=2
#SBATCH --ntasks-per-node=1
#SBATCH --gpus-per-node=4
#SBATCH --cpus-per-task=72
#SBATCH --time=02:00:00
#SBATCH --job-name=ddp-training

# ── Resolve head node for torchrun rendezvous ────────────────────────────
nodes=( $( scontrol show hostnames ${SLURM_JOB_NODELIST} ) )
head_node="${nodes[0]}"
head_node_ip=$(srun --nodes=1 --ntasks=1 -w "${head_node}" hostname -I \
               | awk '{print $1}')
echo "Head node: ${head_node} (${head_node_ip})"

# ── Environment: set explicitly so the job does not depend on your submit shell ──
module load python/miniforge3_pytorch/2.12.0
conda activate base

# ── Surface NCCL errors in the job log (optional) ─────────────────────────
export NCCL_DEBUG=WARN

# ── Launch DDP ────────────────────────────────────────────────────────────
srun torchrun \
     --nnodes=${SLURM_NNODES} \
     --nproc_per_node=${SLURM_GPUS_PER_NODE} \
     --rdzv_id=$RANDOM \
     --rdzv_backend=c10d \
     --rdzv_endpoint="${head_node_ip}:29500" \
     my_ddp_training.py

Important

Do not set NCCL_SOCKET_IFNAME, FI_PROVIDER, NCCL_NET_OFI_PROVIDER, or the other NCCL_* / FI_CXI_* tuning variables yourself. The nccl-ofi-plugin module in the default login environment already sets a tuned Slingshot surface – including NCCL_SOCKET_IFNAME=hsn – and a hand-copied value from another cluster’s example overrides it and can push traffic onto the management Ethernet. Override one only to test a specific symptom, and unset it afterwards. See the NCCL section in the programming environment guide for the full list.

Tip

For GPU-aware MPI workflows (passing device tensors directly to MPI), the module sets MPICH_GPU_SUPPORT_ENABLED=1 automatically.

See PyTorch Multi-Node in the running jobs guide for a complete worked example with collapsible script.

NGC Containers

NVIDIA NGC containers are available at /sw/user/NGC_containers/. List available PyTorch containers:

ls /sw/user/NGC_containers/pytorch_*

The most recent container is pytorch_26.01-py3.sif.

#!/bin/bash
#SBATCH --account=account_name
#SBATCH --partition=ghx4
#SBATCH --nodes=1
#SBATCH --gpus-per-node=1
#SBATCH --time=00:20:00

srun apptainer run --nv \
     --bind /projects \
     /sw/user/NGC_containers/pytorch_26.01-py3.sif \
     python3 my_script.py

Note

NGC containers bundle their own CUDA runtime and libraries. For multi-node container jobs, additional bind mounts for Cray MPICH and Slingshot libraries are required – see Containers.

Installing PyTorch in Your Own Environment

If you need a PyTorch version that is not available as a module, you can install into your own conda or venv environment. Use the stable release index, not the nightly index.

# Create a fresh conda environment
module load python/miniforge3_pytorch       # any version -- this only provides conda
conda create -p ~/my_pytorch python=3.12 -y
conda activate ~/my_pytorch

# Install PyTorch with CUDA 13.0 support (aarch64 wheels available)
pip install torch torchvision torchaudio \
    --index-url https://download.pytorch.org/whl/cu130

To find wheel availability for other CUDA versions or older releases, check the PyTorch previous versions page. Not every CUDA variant publishes aarch64 wheels.

If your custom environment needs multi-node MPI support, build mpi4py from source against Cray MPICH:

MPICC="cc -shared" pip install --no-cache-dir --no-binary mpi4py mpi4py

Confirm it linked correctly before running a multi-node job (run on a compute node):

ldd $(python -c "import mpi4py.MPI as m; print(m.__file__)") | grep -E "mpi|fabric"

Expect libmpi_gnu_123.so.12 and libmpi_gtl_cuda.so.0 resolving from /opt/cray/pe/lib64/, and libfabric.so.1 from /opt/cray/libfabric/. A generic libmpi.so, or a libfabric resolving inside your environment, means the build did not pick up Cray MPICH.

Warning

Never conda install mpi4py. The conda package links against its own MPI and pulls in a conda libfabric that shadows the Slingshot one. Neither substitution reports an error – inter-node traffic silently degrades to a TCP fallback.

Building Custom CUDA Extensions

torch.utils.cpp_extension compiles custom CUDA kernels against whatever CUDA_HOME points at, and CUDA_HOME is always set on DeltaAI – by the default cudatoolkit module under python/miniforge3_pytorch/2.10.0, and by the site cuda/13.1.1 module that 2.11.0 and 2.12.0 load themselves. Check which toolkit you are building against before you compile:

module load python/miniforge3_pytorch/2.12.0
conda activate base
echo $CUDA_HOME        # points to the system CUDA toolkit
python -c "from torch.utils.cpp_extension import CUDAExtension; print('OK')"

Note

Under the 2.11.0 and 2.12.0 modules, CUDA_HOME resolves to /sw/user/cudatoolkits/installs/cuda-13.1.1 while the PyTorch wheel itself bundles the CUDA 13.0 runtime, so an extension you build is compiled against a slightly newer toolkit than the one torch ships. Record both versions alongside any extension you distribute to other users, and rebuild the extension when you move it to a different PyTorch module.