How do I scale distributed training across multiple GPUs?

To scale distributed training across multiple GPUs, you usually use data parallelism or distributed data parallelism (DDP). The basic idea is:

1. Replicate the model on each GPU.
2. Split the input batch across GPUs.
3. Each GPU computes gradients on its mini-batch.
4. Synchronize gradients across GPUs.
5. Update model weights consistently.

Most common approach: Distributed Data Parallel (recommended)

For most frameworks, DDP is the standard way to scale training efficiently.

Why DDP?

• Much faster and more stable than older single-process multi-GPU approaches.
• Scales well across multiple GPUs and multiple machines.
• Each GPU gets its own process, reducing contention.

Typical workflow

• Launch one process per GPU.
• Initialize a process group.
• Wrap the model in DDP.
• Use a distributed sampler so each process sees different data.
• Aggregate gradients automatically during backprop.

Example in PyTorch

import torch
import torch.distributed as dist
import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler

def setup(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)

def cleanup():
    dist.destroy_process_group()

def train(rank, world_size):
    setup(rank, world_size)

    model = nn.Linear(10, 1).to(rank)
    ddp_model = DDP(model, device_ids=[rank])

    dataset = ...
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
    loader = DataLoader(dataset, batch_size=32, sampler=sampler)

    optimizer = torch.optim.Adam(ddp_model.parameters(), lr=1e-3)
    loss_fn = nn.MSELoss()

    for epoch in range(10):
        sampler.set_epoch(epoch)
        for x, y in loader:
            x, y = x.to(rank), y.to(rank)
            optimizer.zero_grad()
            loss = loss_fn(ddp_model(x), y)
            loss.backward()
            optimizer.step()

    cleanup()

Launch with:

torchrun --nproc_per_node=4 train.py

Key things to watch

1. Batch size

• Usually you increase the global batch size as you add GPUs.
• Global batch size = per-GPU batch size × number of GPUs.
• You may need to adjust learning rate accordingly.

2. Data loading

• Use a DistributedSampler so each GPU processes unique data.
• Set the epoch on the sampler each epoch for proper shuffling.

3. Communication overhead

• Gradient synchronization can become a bottleneck.
• Use fast interconnects like NVLink, InfiniBand, or Ethernet at scale.
• Mixed precision can help reduce bandwidth and memory use.

4. Multi-node training

• Works the same way, but you must set up networking, ranks, and world size correctly.
• Make sure firewall/network config allows process-group communication.

Other scaling strategies

• Model parallelism: split the model across GPUs when it doesn’t fit on one GPU.
• Pipeline parallelism: divide model into stages and process micro-batches.
• Tensor parallelism: split large matrix operations across GPUs.
• FSDP / ZeRO: shard model states, gradients, and optimizer states to save memory.

Practical recommendations

• Start with DDP if your model fits on one GPU.
• Use mixed precision for speed and memory savings.
• Profile training to find bottlenecks in input pipeline vs communication.
• For very large models, consider FSDP or DeepSpeed ZeRO.

If you want, I can also show:

• a PyTorch multi-GPU training template
• a TensorFlow distribution strategy example
• or how to scale on multiple machines.