Multi-GPU and Distributed Computing

torch.nn.DataParallel is the simplest multi-GPU wrapper:

model = nn.DataParallel(model)  # wraps model for multi-GPU
output = model(input)           # splits input across GPUs

How it works:

1. Replicates model to each GPU
2. Splits the input batch across GPUs
3. Forward pass runs in parallel on each GPU
4. Outputs are gathered to GPU 0
5. Loss and backward run on GPU 0
6. Gradients are broadcast to update all replicas

Limitations: GPU 0 is a bottleneck (gathers all outputs, computes loss); uses Python threads (GIL-limited); does not scale beyond a single machine.

DistributedDataParallel is the production-grade multi-GPU solution:

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

dist.init_process_group('nccl')
local_rank = int(os.environ['LOCAL_RANK'])
model = model.to(local_rank)
model = DDP(model, device_ids=[local_rank])

Key advantages over DataParallel:

• One process per GPU (no GIL contention)
• Gradient all-reduce overlaps with backward pass
• Scales to multiple machines via NCCL
• No GPU 0 bottleneck (symmetric across all GPUs)

All-reduce is the fundamental operation in distributed training. It computes the element-wise sum (or average) of tensors across all GPUs and distributes the result to every GPU:

$$\mathbf{g}_{\text{avg}} = \frac{1}{N_{\text{GPU}}} \sum_{i=1}^{N_{\text{GPU}}} \mathbf{g}_i$$

where $\mathbf{g}_i$ is the gradient on GPU $i$.

# Manual all-reduce
dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
tensor /= dist.get_world_size()

DDP performs all-reduce automatically during backward(), overlapping communication with gradient computation using bucketed all-reduce.

When a model does not fit on a single GPU, split it across GPUs:

Pipeline parallelism: Split layers across GPUs.

class PipelineModel(nn.Module):
    def __init__(self):
        self.part1 = nn.Sequential(...).to('cuda:0')
        self.part2 = nn.Sequential(...).to('cuda:1')

    def forward(self, x):
        x = self.part1(x.to('cuda:0'))
        x = self.part2(x.to('cuda:1'))
        return x

Tensor parallelism: Split individual layers (e.g., split a large linear layer's weight matrix column-wise across GPUs).

FSDP (Fully Sharded Data Parallel): Shards parameters, gradients, and optimizer states across GPUs, gathering them on-demand.

PyTorch supports three communication backends:

# Initialize with NCCL (best for GPU)
dist.init_process_group(
    backend='nccl',
    init_method='env://',
)

NCCL (NVIDIA Collective Communications Library) provides optimized GPU-to-GPU communication that bypasses the CPU entirely via NVLink or PCIe peer-to-peer.