Multi-GPU Programming

Scaling Beyond One GPU

Large models don't fit on a single GPU. Even when they do, training is too slow. Multi-GPU training addresses both problems, but introduces communication overhead.

The Scale Problem

The LLaMA paper trained 65B parameters on 2048 A100 GPUs for 21 days. Scaling efficiently across thousands of GPUs is essential for frontier model training.

Scaling Efficiency

Perfect scaling means N GPUs give N× throughput. Reality is worse due to:

• Communication overhead — time spent synchronizing gradients
• Bubble time — GPUs waiting for data from other GPUs
• Memory overhead — duplicate buffers, communication buffers
• Load imbalance — some GPUs finish before others

Data, Tensor, and Pipeline

There are three fundamental ways to distribute work across GPUs. Modern large-scale training combines all three in what's called 3D parallelism.

When to Use Each

Collective Communication Operations

NCCL (NVIDIA Collective Communications Library) provides optimized primitives for GPU-to-GPU communication. Understanding these operations is essential for distributed training.

AllReduce

The most common operation: sum gradients across all GPUs, result available on all GPUs.

AllGather

Concatenate data from all GPUs onto all GPUs. Used in tensor parallelism to reconstruct tensors.

ReduceScatter

Reduce then scatter: each GPU gets a different portion of the reduced result. This is the key operation in ZeRO Stage 2.

Communication Cost

NVLink, NVSwitch, and Interconnects

Communication bandwidth determines which parallelism strategies are practical. GPU interconnect technology varies dramatically in bandwidth.

Interconnect Bandwidth

DGX H100 Topology

The DGX H100 uses NVSwitch to provide full-bisection bandwidth between all 8 GPUs:

Multi-Node Considerations

Inter-node bandwidth (InfiniBand/Ethernet) is 10-50× slower than NVLink. This constrains which parallelism works across nodes:

• Data parallelism — works well (gradient sync once per step)
• Pipeline parallelism — works well (point-to-point only)
• Tensor parallelism — avoid across nodes (communication every layer)

Implementing DDP

DistributedDataParallel (DDP) is the workhorse of multi-GPU training. Each GPU has a full model copy, processes different data, and synchronizes gradients.

Basic DDP Setup

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

# Initialize process group (called once per process)
dist.init_process_group(backend="nccl")

# Get local rank (which GPU this process uses)
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)

# Create model and wrap with DDP
model = MyModel().cuda()
model = DDP(model, device_ids=[local_rank])

# Create distributed sampler for data
train_sampler = DistributedSampler(train_dataset)
train_loader = DataLoader(
    train_dataset,
    sampler=train_sampler,
    batch_size=per_gpu_batch_size
)

# Training loop (gradients auto-synced by DDP)
for epoch in range(epochs):
    train_sampler.set_epoch(epoch)  # Shuffle per epoch
    for batch in train_loader:
        loss = model(batch)
        loss.backward()  # DDP syncs gradients here
        optimizer.step()
        optimizer.zero_grad()

How DDP Works

1. Bucket gradients — group gradients into ~25MB buckets for efficient AllReduce
2. Overlap compute and communication — AllReduce starts as soon as bucket is ready
3. Broadcast initial weights — rank 0 broadcasts weights to all ranks at start

# Single node, 8 GPUs
torchrun --nproc_per_node=8 train.py

# Multi-node (run on each node)
torchrun --nnodes=2 --nproc_per_node=8 \
    --rdzv_id=job1 --rdzv_backend=c10d \
    --rdzv_endpoint=master_ip:29500 train.py

Scaling Efficiency

DDP efficiency depends on the ratio of compute to communication:

Memory-Efficient Data Parallelism

ZeRO (Zero Redundancy Optimizer) eliminates memory redundancy in data parallelism. Instead of each GPU storing full optimizer states, gradients, and parameters—they're partitioned.

ZeRO Stages

Baseline DDP — each GPU stores everything:

ZeRO Stage 1 — partition optimizer states:

ZeRO Stage 2 — partition gradients (ReduceScatter instead of AllReduce):

ZeRO Stage 3 — partition parameters (AllGather before each forward/backward):

Memory Savings

Ψ = model parameters. Adam FP32 optimizer states = 12 bytes/param (momentum + variance + master weight).

import deepspeed

# Config enables ZeRO Stage 2
ds_config = {
    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {"device": "cpu"},  # Optional CPU offload
        "contiguous_gradients": True,
        "overlap_comm": True
    },
    "fp16": {"enabled": True},
    "train_batch_size": 32
}

model, optimizer, _, _ = deepspeed.initialize(
    model=model,
    config=ds_config
)

Splitting Layers Across GPUs

Megatron-LM style tensor parallelism splits individual linear layers across GPUs. This is essential for models where a single layer doesn't fit in GPU memory.

Column-Parallel Linear

Split the weight matrix along the output dimension. Each GPU computes part of the output.

# Weight W shape: (in_features, out_features)
# Split into W = [W_0, W_1] along columns

# GPU 0: A_0 = X @ W_0   (all of X, part of W)
# GPU 1: A_1 = X @ W_1   (all of X, part of W)

# Result: A = [A_0, A_1] (concatenated)
# Requires AllGather to get full A on all GPUs (if needed)

Row-Parallel Linear

Split the weight matrix along the input dimension. Requires input to be partitioned.

# Weight W shape: (in_features, out_features)
# Split into W = [W_0; W_1] along rows
# Input A = [A_0, A_1] partitioned across GPUs

# GPU 0: B_0 = A_0 @ W_0  (partial product)
# GPU 1: B_1 = A_1 @ W_1  (partial product)

# Result: B = B_0 + B_1 (requires AllReduce)

Transformer Tensor Parallelism

In transformers, column-parallel and row-parallel layers are combined to minimize communication:

This pattern requires only 2 AllReduces per transformer layer (one after attention, one after FFN).

Key Takeaways