The Ultimate Guide to Installing PyTorch on a GPU Server

In this tutorial, we will walk you through the exact steps to install NVIDIA drivers, the CUDA toolkit, cuDNN, and PyTorch on a GPU dedicated server running Ubuntu.

Prerequisites

Before we begin, ensure you have the following:

• Hardware: A GPU Dedicated Server equipped with an NVIDIA GPU (e.g., RTX 3090, RTX 4090, RTX 5080, or enterprise equivalents).
• Operating System: A fresh installation of Ubuntu 22.04 LTS.
• Access: Root or sudo user access to your server via SSH.

Verify you have sudo access by logging in and running:

sudo whoami
# Should output: root

Step 1: Clean Up & Update Your System

First, ensure your system packages are up to date. To avoid package conflicts, it is a best practice to purge any existing, pre-installed NVIDIA drivers before starting fresh.

sudo apt update && sudo apt upgrade -y
sudo apt remove --purge '^nvidia-.*' -y
sudo apt autoremove -y
sudo reboot

Step 2: Install NVIDIA Drivers

The NVIDIA driver is the foundation of your GPU software stack. Once your server reboots, reconnect via SSH and add the official NVIDIA driver PPA to install the latest stable production driver (we will use version 560 as our baseline).

sudo add-apt-repository ppa:graphics-drivers/ppa -y
sudo apt update
sudo apt install nvidia-driver-560 -y
sudo reboot

After the second reboot, verify the driver is loaded correctly by running the System Management Interface tool:

nvidia-smi

You should see a table displaying your GPU model, driver version, and CUDA version.

Step 3: Install the CUDA Toolkit and cuDNN

While PyTorch ships with its own CUDA runtime, having the full CUDA Toolkit installed on your system is necessary for compiling custom CUDA kernels and libraries like FlashAttention.

Install CUDA 12.4

(Highly stable for modern PyTorch builds):

wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb
sudo dpkg -i cuda-keyring_1.1-1_all.deb
sudo apt update
sudo apt install cuda-toolkit-12-4 -y

Add CUDA to your system PATH:

echo 'export PATH=/usr/local/cuda-12.4/bin:$PATH' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda-12.4/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc
source ~/.bashrc

Next, install cuDNN 9. cuDNN drastically accelerates deep learning primitives like convolutions and attention layers:

sudo apt install cudnn9-cuda-12 -y

Step 4: Set Up Conda & Create a Virtual Environment

Using isolated Python environments prevents dependency conflicts between different AI projects. We will use Miniconda to create a clean environment.

Download and initialize Miniconda:

wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
bash Miniconda3-latest-Linux-x86_64.sh -b -p $HOME/miniconda3
eval "$($HOME/miniconda3/bin/conda shell.bash hook)"
conda init bash
source ~/.bashrc

Create and activate a dedicated Python 3.11 environment named pytorch_env:

conda create -n pytorch_env python=3.11 -y
conda activate pytorch_env

(You should now see (pytorch_env) at the beginning of your terminal prompt).

Step 5: Install PyTorch with GPU Support

With your environment active, use pip to install the CUDA 12.4-compatible build of PyTorch. (Using pip is generally recommended over conda for fetching the latest official PyTorch wheels).

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124

Step 6: Verify the Installation and Benchmark Compute

It is critical to test that PyTorch recognizes your GPU and is ready for heavy workloads. Copy and paste the following command into your terminal to generate and run a benchmarking script. It will check your system diagnostics and run a quick Matrix Multiplication (TFLOPS) benchmark:

cat << 'EOF' > benchmark.py
import torch
import time

print(f'PyTorch version: {torch.__version__}')
print(f'CUDA available:  {torch.cuda.is_available()}')
print(f'CUDA version:    {torch.version.cuda}')
print(f'cuDNN version:   {torch.backends.cudnn.version()}')
print(f'GPU count:       {torch.cuda.device_count()}')
print(f'GPU name:        {torch.cuda.get_device_name(0)}')

print('\nRunning Matrix Multiplication Benchmark...')
size = 8192
a = torch.randn(size, size, device='cuda')
b = torch.randn(size, size, device='cuda')
torch.cuda.synchronize()

start = time.time()
for _ in range(10):
    c = torch.matmul(a, b)
torch.cuda.synchronize()
elapsed = time.time() - start

tflops = (2 * size**3 * 10) / elapsed / 1e12
print(f'Matrix multiply: {elapsed:.2f}s, ~{tflops:.1f} TFLOPS')
EOF

python benchmark.py

Expected Output (Example):

PyTorch version: 2.x.x+cu124
CUDA available:  True
CUDA version:    12.4
cuDNN version:   90100
GPU count:       1
GPU name:        NVIDIA GeForce RTX 4090

Running Matrix Multiplication Benchmark...
Matrix multiply: 0.24s, ~45.8 TFLOPS

Conclusion

Congratulations! You have successfully configured a production-ready software stack on your GPU dedicated server. Your system is now running isolated, optimized, and fully accelerated PyTorch.