Multi-GPU from First Principles

Adding a GPU

Additional GPUs provide both compute and memory, subject to the following constraint: communication has a measurable cost, especially on a machine without NVLink.

Plan: build data parallelism by hand, derive the collective the professionals use, then measure what a second GPU costs — and predict, before running, whether it pays.

Three Ways to Split

Data parallel is our subject: simplest, one sync per step, works for any model that fits. Pipeline and tensor parallel wait for the Language Models part.

Data Parallelism by Hand

Split batch → forward/backward per replica → allreduce gradients → identical update. One process; tensors moved explicitly.

def train_batch(X, y, device_params, devices, lr):
    X_shards, y_shards = split_batch(X, y, devices)
    ls = [loss(lenet(Xs, dev_W), ys).sum()
          for Xs, ys, dev_W in zip(X_shards, y_shards, device_params)]
    for l in ls:
        l.backward()
    with torch.no_grad():
        for i in range(len(device_params[0])):
            allreduce([device_params[c][i].grad for c in range(len(devices))])
        for param in device_params:
            d2l.sgd(param, lr, X.shape[0])
            for p in param:
                p.grad = None

Two GPUs, No Speedup

train(num_gpus=min(2, d2l.num_gpus()), batch_size=256, lr=0.2)

test acc 0.79, 1.85 sec/epoch on 2 GPU(s)

LeNet is too small for effective data parallelism: halving a small batch underutilizes each GPU. Not a bandwidth problem — the whole gradient set is about half a megabyte. Wrong regime, not wrong technique.

Ring Allreduce

Star: hub moves (k-1)N. Ring (reduce-scatter + all-gather): \frac{2(k-1)}{k}N per device — nearly constant, bounded by 2N for any k. The identity that becomes FSDP.

Communication Cost and Scaling

t_{\text{step}}(k) \approx t_{\text{compute}}(B/k) + 2N/\beta

if d2l.num_gpus() >= 2:
    N = 64 * 1024 * 1024  # 64M floats = 256 MB per replica
    data = [torch.randn(N, device=d2l.try_gpu(i)) for i in range(2)]
    t = d2l.Benchmark(lambda: allreduce(data), warmup=2, repeats=5).time
    # Per-device traffic is ~2N bytes; report effective bandwidth
    print(f'allreduce {2 * N * 4 / t / 1e9:.2f} GB/s effective '
          f'over {1000 * t:.1f} ms')
else:
    print('needs 2 GPUs')

allreduce 17.68 GB/s effective over 30.4 ms

Raw copies sustain tens of GB/s (PCIe-limited); NCCL’s fallback transport achieves lower bandwidth because one stage limits it; §13.6 measures the env-switch workaround, and its limits. So LeNet’s no-speedup isn’t communication — it’s t_{\text{compute}}(B/k) not shrinking when a small batch is halved. A larger model and batch can use the second GPU effectively, as the next section demonstrates.

Lineage

• Parameter servers (push/pull): the asynchronous, multi-machine era; alive in recsys embeddings and other sparse, asynchronous state.
• Synchronous collectives (ring allreduce): won for dense training; what DDP runs.

Production map → the Tools appendix. Next: let the library run the ring for us.