PyTorch

PyTorch is a popular python deep learning/autodifferentiation/optimization library that has excellent GPU and CPU support. It features flexible eager mode execution, just-in-time compilation (“JIT”) support, and support for domain-specific tools (e.g., torchvision for image-based learning tasks). It can be loaded in a python environment, and the presence of GPU accelerators can be tested as such:

Python 3.10.9 (main, Jan 11 2023, 15:21:40) [GCC 11.2.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>>
>>> import torch
>>> for i in range(torch.cuda.device_count()):
...    print(torch.cuda.get_device_properties(i).name)
...
NVIDIA A100-SXM4-80GB
NVIDIA A100-SXM4-80GB
NVIDIA A100-SXM4-80GB
NVIDIA A100-SXM4-80GB

Extensions

The anaconda3/2023.09 module’s python distribution also contains some useful extensions to PyTorch :

PyTorch Lightning - Powerful, HPC-friendly, boilerplate-removing library for training, logging, and reproducibility with deep learning models.
PyTorch Geometric - Flexible graph neural network package for use in molecular/materials science, network science, and many other application domains of graph theory.

Examples

Examples of CPU, (multi) GPU, and multi-node training tasks for HPC environments can be found here. Below are reproduced examples for training convolutional neural network image classification models on the Fashion-MNIST dataset.

Setup (on login node):

This sets up some simple packages:

$ module load anaconda3/2023.09
$ conda activate base
$ git clone https://github.com/Ruunyox/pytorch-hpc
$ cd pytorch-hpc
$ pip install --user .

1. Single node, single GPU:

We start with a training YAML file (fashion_mnist_conv_gpu.yaml) appropriate for PyTorch Lightning (note that a similar training jobs can be set up without PyTorch Lightning - see the official PyTorch tutorials for more granular examples):

Since only 1 GPU is needed, it is better to use the gpu-a100:shared partition and request just one GPU (gres=gpu:A100:1) rather than queuing for a full node with 4 GPUs. The following SLURM submission script details the options:

#! /bin/bash
#SBATCH -J pyt_cli_test_conv_gpu
#SBATCH -o pyt_cli_test_conv_gpu.out
#SBATCH --time=00:30:00
#SBATCH --partition=gpu-a100:shared
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --gres=gpu:A100:1
#SBATCH --mem-per-cpu=1G
#SBATCH --cpus-per-task=4

module load cuda/11.8
module load anaconda3/2023.09

conda activate base

srun pythpc --config fashion_mnist_conv_gpu.yaml fit

and can be run using:

The results can be inspected using TensorBoard package (also included in the anaconda3/2023.09 module):

which can be viewed on your local machine via SSH tunneling:

Note: you may change the port 8877 to something else if needed. Alternatively, you may copy your events* logfiles to your local machine and inspect them with tensorboard there.

2. Single node, multiple GPUs

Adding more GPUs with Pytorch Lightning is as simple as setting:

In the training yaml (see fashion_mnist_conv_multi_gpu.yaml), and requesting a non-shared partition in the SBATCH options:

Remember that the number of nodes/GPUs requested through SLURM must match those requested in the PyTorch Lightning training YAML.

3. Multiple nodes, multiple GPUs

Training across multiple nodes with multiple GPUs on a cluster is seamless with Pytorch Lightning. Simply change the training YAML to include:

Which expects 2 nodes with 4 GPUs each, for a total of 8 GPUs, using a distributed data parallel strategy (see here for alternative PyTorch Lightning distributed training strategies). Accordingly, the SLURM submission script must now be changed to include: