PyTorch is a machine learning library based on the Torch library, used for applications such as computer vision and natural language processing, originally developed by Meta AI and now part of the Linux Foundation umbrella.

PyTorch provides two high-level features:

• Tensor computing (like NumPy) with acceleration via GPUs
• Deep neural networks built on a tape-based automatic differentiation system

DCS (AiMOS) Cluster

Since the official PyTorch distribution does not include PPC64le wheels, running on DCS(AiMOS) requires manually compilation.

• The following directions assume a working Conda install.

• Installation requires proxy access to allow external downloads

Pull PyTorch source code:

git clone https://github.com/pytorch/pytorch
cd pytorch
git checkout v2.6.0
git submodule sync && git submodule update --init --recursive

Create a new Conda environment:

conda create -n torch260 python=3.11
conda activate torch260

NOTE! You may need to tell Conda to use the proper install for the current architecture, see this link for more information on that

Install build tools and CUDA to Conda environment:

NOTE! You may substitute a different CUDA version here, if necessary (eg, cuda-12.1.0 also works)

conda install -y -c conda-forge gcc=12.4.0 gxx=12.4.0
conda install -y -c conda-forge numpy=1.26.4
conda install -y nvidia/label/cuda-12.4.1::cuda
conda install -y conda-forge::libopenblas
conda install -y cmake ninja

Install CUDNN 9.0:

Since the default Conda repos only provide CUDNN up to version 8.9, it is necessary to manually download it from Nvidia and import it into the Conda environment

wget https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-ppc64le/cudnn-linux-ppc64le-9.0.0.312_cuda12-archive.tar.xz

tar xf cudnn-linux-ppc64le-9.0.0.312_cuda12-archive.tar.xz

CUDNN_DIR=</PATH/TO/YOUR/EXTRACTED/CUDNN/DIRECTORY>
rsync -av $CUDNN_DIR/include/ $CONDA_PREFIX/include/
rsync -av $CUDNN_DIR/lib/ $CONDA_PREFIX/lib/

Install PyTorch dependencies:

NOTE! These must be installed one at a time. A single conda install with all packages can time out due to solver complexity.

conda install -c conda-forge -y astunparse
conda install -c conda-forge -y expecttest
conda install -c conda-forge -y hypothesis
conda install -c conda-forge -y psutil
conda install -c conda-forge -y pyyaml
conda install -c conda-forge -y requests
conda install -c conda-forge -y setuptools
conda install -c conda-forge -y types-dataclasses
conda install -c conda-forge -y typing-extensions
conda install -c conda-forge -y sympy
conda install -c conda-forge -y filelock
conda install -c conda-forge -y networkx
conda install -c conda-forge -y jinja2
conda install -c conda-forge -y fsspec
conda install -c conda-forge -y lintrunner
conda install -c conda-forge -y packaging
conda install -c conda-forge -y optree

Create a build job

Create a new sbatch file to submit the compile as a Slurm job.

• Change the "PATH TO YOUR MINICONDA INSTALL" line as appropriate
• Change the 'conda activate' line to match the name set earlier
• Set the value of PYTORCH_BUILD_VERSION to the exact version previously downloaded.

BEWARE! An incorrect version here will still produce a 'working' wheel, but will cause version-check errors downstream with various downstream packages.

build_torch.sh

#!/bin/bash
#SBATCH --job-name=torch                # short name for your job
#SBATCH -n 1                            # request a single node
#SBATCH --partition=dcs-2024            # request a node on the 'dcs-2024' partition
#SBATCH --gres=gpu:1                    # number of allocated gpus per node
#SBATCH --output=./dcs/torch_build.log  # Standard output and error log
#SBATCH --time=06:00:00                 # total run time limit (HH:MM:SS)

eval "$(/<PATH TO YOUR MINICONDA INSTALL>/miniconda3/bin/conda shell.bash hook)" 
conda activate torch260

# Export a few environment variables for the compiler
export CC=$(which gcc)
export CXX=$(which g++)
export CUDA_HOME="$CONDA_PREFIX"
export CUDA_TOOLKIT_ROOT_DIR="$CONDA_PREFIX"
export CUDACXX="$CONDA_PREFIX/bin/nvcc"
export CUDNN_ROOT="$CONDA_PREFIX"
export CUDNN_INCLUDE_DIR="$CONDA_PREFIX/include"
export CUDNN_LIBRARY="$CONDA_PREFIX/lib/libcudnn.so"
export PATH="$CONDA_PREFIX/bin:$PATH"
export LD_LIBRARY_PATH="$CONDA_PREFIX/lib:${LD_LIBRARY_PATH:-}"
export CMAKE_PREFIX_PATH="${CONDA_PREFIX}:${CMAKE_PREFIX_PATH:-}"
export LDFLAGS="-Wl,-rpath,$CONDA_PREFIX/lib"
export CXXFLAGS=""  
export CFLAGS=""

# For debugging - print the libraries in your conda environment
echo "Checking for CUDA libraries in $CONDA_PREFIX/lib:"
ls -l $CONDA_PREFIX/lib/libcudart*
ls -l $CONDA_PREFIX/lib/libcupti*
ls -l $CONDA_PREFIX/lib/libopenblas*
ls -l $CONDA_PREFIX/lib/libgomp*

#Build!
python setup.py clean

export PYTORCH_BUILD_VERSION="2.6.0+cu124"
export PYTORCH_BUILD_NUMBER=0
python setup.py bdist_wheel
echo "Wheel file created in dist/ directory:"
ls -lh dist/*.whl

Then submit with

sbatch build_torch.sh

This will take approximately 90-120 minutes. Output is written to the value defined by the #SBATCH --output option

Install the wheel:

Install the newly built wheel to your Conda environment

pip install dist/torch-2.6.0+cu124-cp311-cp311-linux_ppc64le.whl

Test installation:

To test, Create a new Slurm submit file, to test against a compute node:

verify_submit.sh

#!/bin/bash
#SBATCH --job-name=torch_test
#SBATCH -n 1
#SBATCH --partition=dcs-2024     
#SBATCH --gres=gpu:1             
#SBATCH --output=./dcs/torch_build.log  
#SBATCH --time=06:00:00          

eval "$(/<PATH TO YOUR MINICONDA INSTALL>/miniconda3/bin/conda shell.bash hook)" 
conda activate torch260

python dcs/verify_pytorch.py

NOTE! Remember to change your Miniconda path and set the proper Conda environment name

Then, create a verify_pytorch.py file with the following contents.

verify_pytorch.py

import sys
import platform
import torch
import numpy as np

def print_separator():
    print("-" * 80)

def check_version():
    print_separator()
    print(f"Python version: {platform.python_version()}")
    print(f"PyTorch version: {torch.__version__}")
    print(f"NumPy version: {np.__version__}")
    print_separator()

def check_cuda():
    print("CUDA availability:")
    print(f"CUDA available: {torch.cuda.is_available()}")

    if torch.cuda.is_available():
        print(f"CUDA version: {torch.version.cuda}")
        print(f"cuDNN version: {torch.backends.cudnn.version() if torch.backends.cudnn.is_available() else 'Not available'}")
        print(f"Number of CUDA devices: {torch.cuda.device_count()}")

        for i in range(torch.cuda.device_count()):
            print(f"Device {i}: {torch.cuda.get_device_name(i)}")
            print(f"  Compute capability: {torch.cuda.get_device_capability(i)}")
            print(f"  Memory: {torch.cuda.get_device_properties(i).total_memory / (1024 ** 3):.2f} GB")
    print_separator()

def test_cpu_tensor():
    print("Testing CPU tensor operations:")
    try:
        # Create and manipulate CPU tensors
        x = torch.randn(3, 3)
        y = torch.randn(3, 3)
        z = x @ y  # Matrix multiplication
        print("CPU tensor operations: SUCCESS")
        print(f"Sample tensor calculation result:\n{z}")
    except Exception as e:
        print(f"CPU tensor operations: FAILED - {str(e)}")
    print_separator()

def test_cuda_tensor():
    print("Testing CUDA tensor operations:")
    if not torch.cuda.is_available():
        print("CUDA tensor operations: SKIPPED (CUDA not available)")
        return

    try:
        # Create and manipulate CUDA tensors
        x = torch.randn(3, 3).cuda()
        y = torch.randn(3, 3).cuda()
        z = x @ y  # Matrix multiplication
        print("CUDA tensor operations: SUCCESS")
        print(f"Sample tensor calculation result:\n{z}")
    except Exception as e:
        print(f"CUDA tensor operations: FAILED - {str(e)}")
    print_separator()

def test_simple_nn():
    print("Testing simple neural network:")
    try:
        # Define a simple neural network
        class SimpleNN(torch.nn.Module):
            def __init__(self):
                super(SimpleNN, self).__init__()
                self.fc1 = torch.nn.Linear(10, 5)
                self.fc2 = torch.nn.Linear(5, 2)

            def forward(self, x):
                x = torch.relu(self.fc1(x))
                x = self.fc2(x)
                return x

        # Create model and test forward pass
        model = SimpleNN()
        x = torch.randn(2, 10)
        output = model(x)

        print("Neural network forward pass: SUCCESS")
        print(f"Input shape: {x.shape}")
        print(f"Output shape: {output.shape}")
        print(f"Output: {output}")

        # Test with CUDA if available
        if torch.cuda.is_available():
            model = model.cuda()
            x = x.cuda()
            output = model(x)
            print("Neural network on CUDA: SUCCESS")
    except Exception as e:
        print(f"Neural network test: FAILED - {str(e)}")
    print_separator()

def test_autograd():
    print("Testing autograd functionality:")
    try:
        x = torch.randn(3, requires_grad=True)
        y = x * 2
        z = y.mean()
        z.backward()
        print("Autograd test: SUCCESS")
        print(f"Gradient: {x.grad}")
    except Exception as e:
        print(f"Autograd test: FAILED - {str(e)}")
    print_separator()

def main():
    print("PyTorch Installation Verification")
    print("=" * 80)

    check_version()
    check_cuda()
    test_cpu_tensor()
    test_cuda_tensor()
    test_simple_nn()
    test_autograd()

    print("Verification complete!")

if __name__ == "__main__":
    main()

verify_pytorch.py will check installed PyTorch and CUDA versions, test CPU/GPU tensor operations, run a small testing neural network and validate autograd.

A working install will output something like this:

Known Issues

This PyTorch build is expected to work on all DCS nodes. However, issues sometimes appear randomly on certain nodes.

Issue Summary

Issue	Symptom	Solution
Missing CUDA driver	torch.cuda.is_available() returns False	Exclude the node or request a different one
Stale GPU memory	CUDA errors at job start	Exclude known problematic nodes (also make helpdesk ticket)
NCCL / InfiniBand memory error	NCCL WARN Call to ibv_reg_mr failed with error Cannot allocate memory	Add --exclusive to the SBATCH file

• Thanks to Yunshi 'Randy' Wen for authoring this documentation

NPL (AiMOSx) Cluster

As NPL is an x86_64 architecture cluster, prebuilt PyTorch wheels are available in the regular Conda repositories.

Setup Environment:

conda create -n "my_pytorch_environment" python=3.10.13

conda activate my_pytorch_environment

Install PyTorch:

module load gcc 
module load cuda/12.1
conda install pytorch=2.4.1

Troubleshooting

• It is no longer necessary to specify a CUDA by installing "pytorch::pytorch-cuda"

• If CUDA is not loaded into Pytorch, performance will suffer

Confirm CUDA is enabled:

python

import torch
torch.cuda.is_available()
true