8x A100 GPU Setup

You've got a massive model ready to fine-tune, but your 8x A100 nodes sit idle because you don't know how to piece together the drivers, containers, and orchestration that make them sing. Hours of training time — and your GPU budget — evaporate while you wrestle with nvidia-smi and NCCL errors. This lesson walks you through setting up an 8x A100 GPU training environment for LLM fine-tuning so that from the first command to the final epoch, your hardware runs at peak efficiency.

The problem this lesson solves

Fine-tuning a modern LLM like Llama 3 70B is impossible on a single consumer GPU. You need the combined memory and compute of multiple A100s — often 8 of them in one node — working as a single unit. But throwing 8 GPUs into a server doesn't give you a training environment. You need:

• Correct drivers and toolkit versions that match your CUDA runtime.
• Orchestration so PyTorch can split your model across GPUs and synchronize gradients.
• Fiber-optic-fast inter-GPU communication via NVLink and InfiniBand, not just PCIe.

If any of these pieces is misconfigured, you get cryptic errors like NCCL error: unhandled cuda error or silent slowdowns where GPUs idle at 20% utilization. This lesson gives you a battle-tested path to avoid those pitfalls.

Core concept / mental model

Think of your 8x A100 node as a small, high-performance data center. Each GPU is a worker, and the NVIDIA Collective Communications Library (NCCL) is the messenger that lets workers talk to each other. The NVIDIA Container Toolkit gives each worker its own clean room (a container) with the exact CUDA version it needs, isolated from the host OS.

When you launch training, PyTorch Distributed (via torch.distributed) acts as the foreman — it assigns each GPU a rank (0 to 7) and coordinates the work. Gradients are averaged across all 8 GPUs after every step, so the model updates once with a global view. This is called data parallelism — each GPU holds a complete copy of the model but processes different batches of data.

Here's the dependency chain: Hardware (A100) → Driver → CUDA Toolkit → Container Runtime → PyTorch → Distributed Framework. Every layer relies on the one below it. Break any link, and training breaks.

How it works step by step

Setting up an 8x A100 GPU training environment follows a logic sequence that mirrors the dependency stack. Go top-down from hardware to software, and you'll avoid version conflicts.

1. Install the NVIDIA driver — this talks directly to the GPU hardware. For A100s, you need driver version 450.80.02 or later (CUDA 11.0+ support). Use nvidia-smi to confirm all 8 GPUs are visible.

2. Install the NVIDIA Container Toolkit — this lets Docker containers access the GPU. It replaces the older nvidia-docker2 approach and works with Docker's --gpus flag.

3. Choose your container image — start from a PyTorch official image (pytorch/pytorch:2.1.0-cuda12.1-cudnn8-devel) or an NGC container that already has CUDA, cuDNN, and NCCL configured. Don't install CUDA directly on the host; keep it in the container for reproducibility.

4. Verify GPU access inside the container — run nvidia-smi inside the container to confirm all 8 GPUs are visible and the correct driver version is reported.

5. Install Python dependencies — inside the container, install PyTorch with the CUDA version that matches your driver, plus libraries like accelerate, transformers, and numpy.

6. Set up distributed environment variables — PyTorch needs to know the master address, port, and world size. Tools like torchrun handle this automatically, but you need to set NCCL_IB_DISABLE=0 (for InfiniBand) or 1 (if not available).

Hands-on walkthrough

Let's build the environment from scratch on an Ubuntu 22.04 host with 8x A100 GPUs. We'll use Docker for the containerization step because it's the standard in production.

Step 1: Install the NVIDIA driver

# Check if the driver is already installed
nvidia-smi

# If not, install the recommended driver (example for Ubuntu 22.04)
sudo apt update
sudo apt install -y nvidia-driver-535  # or the latest recommended
sudo reboot

# After reboot, verify
expect -c 'spawn nvidia-smi; expect "Driver Version"' || nvidia-smi

Expected output: A table showing 8 NVIDIA A100 GPUs with driver version 535.xx, memory, and current utilization.

Step 2: Install Docker and NVIDIA Container Toolkit

# Install Docker if not present
curl https://get.docker.com | sh
sudo usermod -aG docker $USER
newgrp docker

# Install the NVIDIA Container Toolkit
curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
  sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
  sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list

sudo apt update
sudo apt install -y nvidia-container-toolkit
sudo nvidia-ctk runtime configure --runtime=docker
sudo systemctl restart docker

Step 3: Pull a PyTorch image and test GPU access

docker run --gpus all --rm -it pytorch/pytorch:2.1.0-cuda12.1-cudnn8-devel bash

# Inside the container, check GPUs
nvidia-smi

Expected output: All 8 A100s visible. If you see <none> or fewer GPUs, check your driver and toolkit.

Step 4: Write a minimal distributed training script

Create train.py:

import os
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim
from torch.nn.parallel import DistributedDataParallel as DDP

# Initialize process group
from torch.distributed import init_process_group, destroy_process_group

def setup():
    init_process_group(backend='nccl')
    torch.cuda.set_device(int(os.environ['LOCAL_RANK']))

def cleanup():
    destroy_process_group()

class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(128, 256)

    def forward(self, x):
        return self.fc(x)

def main():
    setup()
    local_rank = int(os.environ['LOCAL_RANK'])
    model = SimpleModel().to(f'cuda:{local_rank}')
    ddp_model = DDP(model, device_ids=[local_rank])

    optim = optim.SGD(ddp_model.parameters(), lr=0.01)
    loss_fn = nn.MSELoss()

    # Fake training data
    data = torch.randn(64, 128).to(f'cuda:{local_rank}')

    for epoch in range(10):
        optim.zero_grad()
        output = ddp_model(data)
        loss = loss_fn(output, torch.randn(64, 256).to(f'cuda:{local_rank}'))
        loss.backward()
        optim.step()
        if dist.get_rank() == 0:
            print(f'Epoch {epoch} loss: {loss.item()}')

    cleanup()

if __name__ == '__main__':
    main()

Run it with torchrun:

torchrun --nproc_per_node=8 train.py

Expected output: Loss decreasing over 10 epochs from all 8 processes (you'll see lines from rank 0).

Compare options / when to choose what

For most ML engineers fine-tuning LLMs, Docker + Container Toolkit is the sweet spot. It gives you reproducibility without the operational burden of Kubernetes.

Pro tip: Always pin your image tags (e.g., pytorch/pytorch:2.1.0-cuda12.1-cudnn8-devel). latest might pull a version that breaks your code.

Troubleshooting & edge cases

nvidia-smi shows no GPUs inside container

• Check that the container runtime is configured: docker info | grep nvidia. If missing, re-run sudo nvidia-ctk runtime configure --runtime=docker.
• Ensure you're using --gpus all (not --runtime=nvidia, which is deprecated).

NCCL error: unhandled cuda error

• This usually means a GPU died or a hang in communication. First, check nvidia-smi for GPU health. Then set NCCL_DEBUG=INFO to get verbose logs. Often it's a driver crash; reboot and test again.
• Also check ulimit -n (file descriptors). Set high: ulimit -n 1048576 in your startup script.

Training is slow: GPUs at 20% utilization

• Check inter-GPU bandwidth: use nvidia-smi topo -m to see if you have NVLink. If not, set NCCL_P2P_DISABLE=1 and rely on network.
• Check if your data loading is the bottleneck. Use DataLoader with num_workers>0 and pin_memory=True.
• Run nsys profile to see where time is spent.

torch.distributed fails to initialize

• Ensure MASTER_ADDR and MASTER_PORT are set correctly. With torchrun, they're auto-set.
• If using InfiniBand, ensure NCCL_IB_DISABLE=0 and the libibverbs libraries are installed. If not, set NCCL_IB_DISABLE=1 to use TCP.

One GPU is slower than others

• This smells like a thermal throttle. Check nvidia-smi for temperature. A100s throttle at 80°C+. Improve airflow or reduce room temperature.
• Also check power caps: nvidia-smi -pl 400 to set a higher power limit if the GPU allows.

What you learned & what's next

You now know how to set up an 8x A100 GPU training environment: install the NVIDIA driver, containerize with the NVIDIA Container Toolkit, verify GPU access, and launch a distributed training script with torchrun. You also learned how to troubleshoot common issues like NCCL errors and slow inter-GPU communication.

To apply this, you can now fine-tune a real LLM. The next lesson in this track covers data loading and streaming for distributed training — you'll learn how to feed large datasets efficiently without bottlenecking your brand-new GPU setup. Your skills from this lesson — especially understanding rank, world size, and NCCL — are the foundation for that. You're one step closer to shipping your own fine-tuned model.

Now go fire up those A100s and experiment with different --nproc_per_node values to see how throughput scales. Happy training!