Installing Pytorch/Pytorch Lightning Using Pip

This guide will walk you through installing Pytorch and/or Pytorch Lighting using Pip. It assumes you have already installed or loaded Miniforge. See the guide on using conda for more. If you haven't loaded Miniforge, you can do so by running

module purge
module load Miniforge3

Setup - Checking Python

If you installed conda on your own and not following our using conda guide, the HPCC may be trying to use the system python installation instead of your own. To test if this is the case, run

which python

If the output starts with /opt/software, you will need to run module unload Python. If the output starts with the path to Miniforge, you don't need to do anything else.

Note

If you are affected by the above issue, you will have to run module unload Python every time you wish to use your own python installation. You may wish to add the module unload Python command to your $HOME/.bashrc file.

Installing Pytorch

Since Pytorch works best when using a GPU, it needs to be installed on a development node with a GPU. We recommend using dev-amd20-v100 for the most compatibility. However, you can also use dev-amd24-h200 for the latest hardware. Run

ssh dev-amd20-v100

Note

You will be restricted to running Pytorch on nodes with v100 GPUs. See the page on cluster resources and SLURM job specifications for more.

You will also need the CUDA compiler, so load this using our module system:

module load CUDA/12.4.0

Note

Please review the Pytorch installation documentation to confirm the CUDA version for the latest stable PyTorch release, along with other dependency information.

It's best practice to use Conda environments to organize your Python packages. To create a new conda environment with the name pytorch run

conda create --name pytorch python pip

Note

We need to include python and pip, otherwise using pip may install directly into the user packages in .local using the system pip installation. This ensures we're using a version of pip installed in the conda environment.

To switch to this new environment, run

conda activate pytorch

Now that you are on a GPU development node, have loaded the CUDA module, and activated your new environment, you can install Pytorch with the following command:

python -m pip install torch torchvision torchaudio

Installing Pytorch Lightning

It's best to install Pytorch following the instructions above before installing Pytorch Lightning, or GPU-support may not function correctly.

After Pytorch has been installed, Pytorch Lightning can be installed to the same pytorch environment using

python -m pip install lightning

Finally, you can confirm Pytorch and Pytorch Lightning have been installed successfully by checking your conda environment for the following packages: torch, torchaudio, torchmetrics, torchvision, pytorch-lightning. Run

conda list | grep torch

In the output you should see all the torch related packages in your conda environment.

# packages in environment at /mnt/home/USERNAME/.conda/envs/pytorch:
pytorch-lightning         2.5.1                    pypi_0    pypi
torch                     2.6.0                    pypi_0    pypi
torchaudio                2.6.0                    pypi_0    pypi
torchmetrics              1.7.1                    pypi_0    pypi
torchvision               0.21.0                   pypi_0    pypi