Building GPT from Scratch — Part 7: Distributed Data Parallel Training

The Problem

Our GPT trains on a MacBook in ~1 minute because it's 3.3M parameters on a small dataset. But real LLMs are different:

• GPT-3: 175B parameters, 300B tokens of training data

• LLaMA 3: 70B parameters, 15T tokens

• Training time on one GPU: years

Even modest models (1B parameters) take weeks on a single A100. The math is simple: if one GPU processes 100K tokens/second and you need to process 1T tokens, that's 115 days.

The solution: split the work across multiple GPUs.

Data Parallelism: The Simplest Multi-GPU Strategy

The idea is straightforward:

1. Replicate the model on every GPU (each GPU has a full copy)

2. Split each training batch into chunks — GPU 0 gets samples 0–15, GPU 1 gets samples 16–31, etc.

3. Compute forward + backward passes independently on each GPU

4. Average the gradients across all GPUs (AllReduce)

5. Update parameters identically on every GPU (same averaged gradients → same update → models stay in sync)

Because every GPU applies the same averaged gradient to the same model state, all copies remain identical after every step. No drift, no divergence.

Why This Works Mathematically

The key insight: gradient of a sum = sum of gradients.

If your batch has 64 samples and you compute the loss as the average across all 64:

L_total = (1/64) * Σ L_i

∂L_total/∂θ = (1/64) * Σ (∂L_i/∂θ)

This sum can be split across GPUs:

GPU 0: compute ∂L/∂θ for samples 0-21   → local_grad_0
GPU 1: compute ∂L/∂θ for samples 22-42  → local_grad_1
GPU 2: compute ∂L/∂θ for samples 43-63  → local_grad_2

Average: (local_grad_0 + local_grad_1 + local_grad_2) / 3 = true gradient

The averaged gradient from 3 GPUs each processing 21 samples is mathematically identical to one GPU processing all 63 samples. Same convergence, 3× throughput.

The AllReduce Operation

AllReduce is the critical communication primitive. It takes a tensor from every GPU and produces the sum (or average) on every GPU:

Before AllReduce:
  GPU 0: grad = [1.0, 2.0, 3.0]
  GPU 1: grad = [4.0, 5.0, 6.0]
  GPU 2: grad = [7.0, 8.0, 9.0]

After AllReduce (sum):
  GPU 0: grad = [12.0, 15.0, 18.0]
  GPU 1: grad = [12.0, 15.0, 18.0]
  GPU 2: grad = [12.0, 15.0, 18.0]

Every GPU ends up with the same result. The implementation uses a ring AllReduce — each GPU sends chunks to its neighbor in a ring pattern. For N GPUs with a gradient of size M:

• Communication volume per GPU: 2M × (N-1)/N

• This is nearly 2M regardless of N — it scales efficiently

For our model's 3.3M parameters (13.2 MB in float32), each AllReduce moves ~26 MB per GPU over NVLink (900 GB/s on H100s) — taking ~0.03ms. Negligible compared to the compute time.

PyTorch DDP: Distributed Data Parallel

PyTorch provides DistributedDataParallel (DDP) — a wrapper that handles all the communication:

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

# Initialize the process group (one process per GPU)
dist.init_process_group(backend="nccl")  # NCCL = NVIDIA's GPU communication library

rank = int(os.environ["LOCAL_RANK"])       # which GPU am I?
world_size = int(os.environ["WORLD_SIZE"]) # how many GPUs total?

# Set this process to use its assigned GPU
torch.cuda.set_device(rank)
device = torch.device(f"cuda:{rank}")

# Wrap the model
model = GPT(config).to(device)
model = DDP(model, device_ids=[rank])

# Training loop is UNCHANGED — DDP hooks into backward() automatically
for step in range(max_steps):
    x, y = get_batch(...)  # each GPU gets different data
    logits, loss = model(x, y)
    loss.backward()        # ← DDP inserts AllReduce here
    optimizer.step()
    optimizer.zero_grad()

What DDP does under the hood:

1. Registers hooks on every parameter's gradient

2. During backward(), as soon as a gradient is computed, it starts the AllReduce for that gradient (overlapping communication with computation)

3. By the time backward() finishes, all gradients are averaged across GPUs

4. optimizer.step() applies the same update everywhere

The overlap is critical — while GPU 0 is still computing gradients for early layers, it's already sending later layers' gradients to other GPUs. This hides most of the communication latency.

Making It Work on Both MacBook and Multi-GPU

The key design: detect the environment and adapt.

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

def setup_distributed():
    """Returns (device, is_distributed, rank, world_size)"""
    world_size = int(os.environ.get("WORLD_SIZE", 1))

    if world_size > 1:
        # Launched with torchrun — multi-GPU mode
        dist.init_process_group(backend="nccl")
        rank = int(os.environ["LOCAL_RANK"])
        torch.cuda.set_device(rank)
        device = torch.device(f"cuda:{rank}")
        return device, True, rank, world_size
    else:
        # Normal launch — single device
        if torch.backends.mps.is_available():
            device = torch.device("mps")
        elif torch.cuda.is_available():
            device = torch.device("cuda")
        else:
            device = torch.device("cpu")
        return device, False, 0, 1

Then the training code:

device, distributed, rank, world_size = setup_distributed()

model = GPT(config).to(device)
if distributed:
    model = DDP(model, device_ids=[rank])

# Only rank 0 prints (avoid duplicate logs from every GPU)
if rank == 0:
    print(f"Training with {world_size} device(s)")

for step in range(max_steps):
    x, y = get_batch(train_dataset, batch_size, device)
    _, loss = model(x, y)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

    if rank == 0 and step % eval_interval == 0:
        print(f"step {step} | loss {loss.item():.4f}")

if distributed:
    dist.destroy_process_group()

Running it:

# On MacBook (single MPS device) — same as before
python train.py

# On a multi-GPU server (4 GPUs)
torchrun --nproc_per_node=4 train.py

# On a cluster (2 nodes × 4 GPUs = 8 total)
torchrun --nnodes=2 --nproc_per_node=4 --master_addr=node0 train.py

Same script. No code branches. The environment determines the behavior.

Data Splitting: DistributedSampler

Each GPU must process different data — otherwise you'd compute the same gradient N times (useless). PyTorch provides DistributedSampler:

from torch.utils.data import DataLoader, DistributedSampler

if distributed:
    sampler = DistributedSampler(
        train_dataset,
        num_replicas=world_size,
        rank=rank,
        shuffle=True,
    )
    dataloader = DataLoader(train_dataset, batch_size=batch_size, sampler=sampler)
else:
    dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

The sampler ensures GPU 0 sees samples [0, 4, 8, 12, ...], GPU 1 sees [1, 5, 9, 13, ...], etc. Every sample is seen exactly once per epoch across all GPUs.

Effective Batch Size

With DDP, the effective batch size = batch_size_per_gpu × world_size.

If each GPU processes 32 samples and you have 4 GPUs, the effective batch is 128. This matters because:

• Larger batches = more stable gradients (less noise)

• But too large = worse generalization (the loss landscape argument)

• Learning rate should scale with batch size (linear scaling rule)

# Scale learning rate with world size
effective_batch_size = batch_size * world_size * grad_accum_steps
base_lr = 3e-4
lr = base_lr * (effective_batch_size / 32)  # baseline was 32

Communication Backends

On MacBook, if you want to simulate DDP for testing:

# Simulates 2 "GPUs" on CPU (slow, but proves the code works)
torchrun --nproc_per_node=2 train.py --device cpu --backend gloo

Scaling Efficiency

DDP scales nearly linearly for models where compute dominates communication:

1 GPU:  1000 tokens/sec
2 GPUs: 1950 tokens/sec  (97.5% efficiency)
4 GPUs: 3800 tokens/sec  (95% efficiency)
8 GPUs: 7400 tokens/sec  (92.5% efficiency)

Efficiency drops slightly as you add GPUs because:

1. AllReduce communication time grows (slightly)

2. GPU idle time waiting for the slowest worker ("straggler effect")

3. At some point, the per-GPU batch becomes too small to saturate compute

For our 3.3M model, the communication is tiny relative to compute — you'd see near-perfect scaling up to dozens of GPUs. For larger models (billions of parameters), the gradient tensors are larger but so is the compute time — the ratio stays favorable.

Going Multi-Node: DDP Across Machines

Everything above works identically across multiple machines. The AllReduce algorithm doesn't care whether GPUs are in the same box or across a datacenter — it just needs a communication channel.

What Changes

Single-node (8 GPUs in one machine):
  Communication: NVLink / PCIe — 900 GB/s / 64 GB/s
  Latency: <1μs
  Launch: torchrun --nproc_per_node=8 train.py

Multi-node (4 machines × 8 GPUs = 32 GPUs):
  Communication: InfiniBand / RoCE — 400 GB/s / 100 GB/s
  Latency: 1-5μs
  Launch: torchrun on each node, coordinated by master

The code is unchanged. Only the launch command and network differ.

How to Launch

Each node runs torchrun with cluster-level arguments:

# On Node 0 (master):
torchrun \
  --nnodes=4 \
  --nproc_per_node=8 \
  --node_rank=0 \
  --master_addr=192.168.1.100 \
  --master_port=29500 \
  train.py

# On Node 1:
torchrun \
  --nnodes=4 \
  --nproc_per_node=8 \
  --node_rank=1 \
  --master_addr=192.168.1.100 \
  --master_port=29500 \
  train.py

# On Node 2, 3: same, with --node_rank=2, --node_rank=3

In practice, a job scheduler (Slurm, Kubernetes) launches all nodes simultaneously:

# Slurm (HPC clusters):
srun --nodes=4 --ntasks-per-node=8 --gpus-per-node=8 \
  torchrun --nproc_per_node=8 train.py

# Kubernetes (cloud):
# A PyTorchJob CRD handles the multi-node coordination

How It Works Under the Hood

1. Rendezvous: all 32 processes connect to the master address (node 0, port 29500). Once all check in, the process group is formed.

2. Rank assignment: each process gets a global rank (0–31) and a local rank (0–7 within its node). LOCAL_RANK determines which GPU on the machine this process uses.

3. NCCL topology detection: NCCL discovers the network — it knows which GPUs are connected by NVLink (same node) vs InfiniBand (cross-node). It automatically builds an optimal communication tree.

4. Hierarchical AllReduce: NCCL optimizes for the topology:

   • First: reduce within each node (fast, NVLink)

   • Then: reduce across nodes (slower, InfiniBand)

   • Finally: broadcast the result back within each node

This minimizes cross-node traffic — only one GPU per node communicates externally.

Network Requirements

For our 3.3M model, the gradient AllReduce is ~13 MB — even 10 Gbps ethernet (1.25 GB/s) handles it in 10ms. Real models with billions of parameters need InfiniBand to keep communication hidden behind compute.

Fault Tolerance

Nodes fail. A 256-GPU training run on 32 machines will hit hardware failures weekly. Options:

1. Checkpoint and restart (simplest):

# Save every N steps
if step % save_interval == 0 and rank == 0:
    torch.save(checkpoint, "checkpoint.pt")

# On restart: all nodes load the same checkpoint
# Lost work = steps since last checkpoint

2. Elastic training (torchrun built-in):

# Allow the job to continue with fewer nodes
torchrun --nnodes=2:4 \       # min 2 nodes, max 4
  --nproc_per_node=8 \
  --rdzv_backend=c10d \       # dynamic rendezvous
  --rdzv_endpoint=master:29500 \
  train.py

If a node dies, remaining nodes re-form the process group and continue. The effective batch size drops (fewer GPUs), but training doesn't stop.

3. In practice: most large-scale training uses periodic checkpointing (every 5–15 minutes) and restarts the full job on failure. The infrastructure cost of occasional lost work is lower than the engineering cost of elastic training at frontier scale.

What DDP Can't Do

DDP requires the entire model to fit on one GPU. Each GPU holds:

• Full model parameters (~4 bytes per param in fp32)

• Gradients (same size as parameters)

• Optimizer states (2× parameters for AdamW — momentum + variance)

Total: ~16 bytes per parameter × 175B parameters = 2.8 TB for GPT-3.

No single GPU holds 2.8 TB. This is where FSDP (Fully Sharded Data Parallel) comes in — the subject of the next post.

Key Takeaways

1. Data parallelism is embarrassingly simple. Split the batch, compute independently, average gradients. Same math, N× throughput.

2. AllReduce is the only communication needed. No complex synchronization, no shared memory, no locks. Just: average this tensor across all GPUs — whether they're in the same box or across a datacenter.

3. DDP overlaps communication with computation. Gradients for later layers ship while earlier layers are still in backward(). Communication is mostly hidden.

4. Same script, different launch. python train.py for single-device, torchrun for multi-GPU, torchrun --nnodes=N for multi-node. Your training loop never changes — NCCL handles the topology.

5. The limit is model size. DDP requires the full model on each GPU. When the model won't fit, you need sharding — which is what FSDP provides.

Next: when the model is too big for one GPU — Fully Sharded Data Parallel (FSDP).

Full code: github.com/the-last-programmers/gpt-from-scratch