Distributed training & systems ML

When a model and its optimiser no longer fit on one GPU, you split them up: copies of the data, slices of each matrix, groups of layers, shards of the optimiser state. This stage is about choosing those splits, the communication each one costs, and how to find out where the time actually goes.

Source: Suraj Sharma, X post (stage 6) · roadmap hub: Module 08

Going deeper: everything below is the full version, with some maths. On a first read, skim the headings and the "Why it matters for robotics" box, and come back for the rest later.

1. What this stage covers

• Learn: 3D parallelism (data, tensor, pipeline), DeepSpeed ZeRO-3, PyTorch FSDP, NCCL communication primitives, gradient checkpointing.
• Practice: shard a 7B model across 4 GPUs, profile the NCCL overhead and tune where the tensor-parallel splits go.
• Why (the post's claim): research at scale needs weights and optimiser states spread across clusters without choking the interconnect. Section 4 checks how far that goes.

2. Concepts, explained

Where the memory goes

Start here, because every technique below attacks one of these terms. With mixed-precision Adam, each parameter costs about 16 bytes before you count activations (ZeRO paper's accounting):

fp16/bf16 weights      2 bytes
fp16/bf16 gradients    2 bytes
fp32 master weights    4 bytes  ┐
fp32 Adam momentum     4 bytes  ├ optimizer state, K = 12
fp32 Adam variance     4 bytes  ┘
total                 16 bytes / parameter  →  7B params ≈ 112 GB

So a full fine-tune of a 7B model does not fit on one 80 GB GPU even before activations. Activations scale with batch × sequence length × hidden size × layers, and at long context they can dominate.

Data parallelism (DP / DDP)

Every GPU holds the full model and processes a different slice of the batch. After the backward pass, gradients are averaged with an all-reduce, so all copies take the same step. PyTorch's DDP buckets gradients and starts all-reducing early buckets while the backward pass is still running, hiding much of the communication. It scales throughput, but memory per GPU doesn't shrink at all.

Tensor parallelism (TP)

Shoeybi, Patwary, Puri, LeGresley, Casper and Catanzaro (2019), Megatron-LM, split each transformer block's matrices so that only a couple of communications are needed:

MLP:   Y = GeLU(X · A) · B
       A split by columns  → each GPU computes GeLU(X · A_i) locally
       B split by rows     → each GPU gets a partial sum Y_i
       all-reduce(Y_i)     → Y           (one all-reduce forward, one backward)
Attention: split by heads, same column-then-row pattern

That's two all-reduces per layer in the forward pass, on activation-sized tensors, every layer. TP therefore needs the fastest links you have, which in practice means NVLink inside one node. They reported 8.3B-parameter models on 512 GPUs at 76% scaling efficiency. Sequence parallelism (Korthikanti et al., 2022) extends TP to split the layer-norm and dropout regions along the sequence, replacing each all-reduce with a reduce-scatter plus all-gather of the same total volume and cutting activation memory.

Pipeline parallelism (PP)

Put consecutive layers on different GPUs and stream micro-batches through them. Day 24 (GPipe) covers the idea and its demo. The cost is the bubble: with p stages and m micro-batches, the idle fraction is about (p − 1)/m. The 1F1B schedule bounds how many micro-batches' activations are in flight. The interleaved schedule from Narayanan et al. (2021) gives each GPU several smaller chunks of layers, cutting the bubble by that factor at the cost of more communication. PP only sends activations between neighbours (point-to-point send/recv), so it tolerates slower links than TP.

3D parallelism

Narayanan et al. composed all three: TP inside a node, PP across nodes, DP over the whole thing. Total GPUs = t × p × d. Their rules of thumb: keep the tensor-parallel degree within a server, use pipeline parallelism to go beyond that, and use data parallelism for the rest. They trained up to a trillion parameters on thousands of GPUs this way. Modern stacks add context parallelism (splitting long sequences, see Stage 03's ring attention) and expert parallelism for MoE.

DeepSpeed ZeRO

Rajbhandari, Rasley, Ruwase and He (2020) noticed that data parallelism stores N identical copies of everything. ZeRO removes the redundancy in three stages. With Ψ parameters and N GPUs:

stage                    per-GPU model-state memory     comm vs DP
DP (no ZeRO)             16Ψ                             1×
ZeRO-1  (P_os)           4Ψ + 12Ψ/N                      1×
ZeRO-2  (P_os+g)         2Ψ + 14Ψ/N                      1×
ZeRO-3  (P_os+g+p)       16Ψ/N                           1.5×

Stages 1 and 2 cost no extra communication, because an all-reduce is already a reduce-scatter followed by an all-gather, and ZeRO just changes what each GPU keeps in between. Stage 3 must all-gather each layer's parameters right before its forward and backward passes, then free them, which adds about 50% more traffic. The paper's example: a 7.5B model on 64 GPUs drops from 120 GB to about 1.9 GB of model state per GPU under stage 3. ZeRO-Offload and ZeRO-Infinity extend this to CPU and NVMe memory.

FSDP

PyTorch's Fully Sharded Data Parallel (Zhao et al., 2023) is the native take on ZeRO-3. You wrap units of the model (usually one per transformer block). Each unit's parameters live sharded, are all-gathered just before use, and are freed afterwards. Gradients are reduce-scattered. Choices that matter: how you wrap (too coarse wastes memory, too fine adds latency), prefetching the next unit's all-gather, mixed-precision policy, and hybrid sharding (shard within a node, replicate across nodes) to keep the heavy traffic on fast links. Newer PyTorch exposes this as fully_shard (FSDP2), built on DTensor, which also composes with tensor parallelism.

NCCL communication primitives

NCCL is NVIDIA's library of GPU collectives. Learn what each does and what moves over the wire for a message of size S on N GPUs (ring algorithms):

all-reduce       everyone ends with the sum        sends ≈ 2(N−1)/N · S per GPU
reduce-scatter   everyone ends with 1/N of sum     sends ≈ (N−1)/N · S
all-gather       everyone ends with all shards     sends ≈ (N−1)/N · S
broadcast        one rank's tensor to all
all-to-all       each rank sends a slice to each   (MoE expert routing)
send / recv      point to point                    (pipeline stages)

The nccl-tests tools report bus bandwidth, which applies these factors so you can compare a measured number directly with the link's hardware speed. Small messages are latency-bound, big ones bandwidth-bound. That's why gradients are bucketed. Debug with NCCL_DEBUG=INFO to see which transport (NVLink, PCIe, InfiniBand, sockets) was actually chosen.

Gradient (activation) checkpointing

Chen, Xu, Zhang and Guestrin (2016): store activations only at some checkpoints and recompute the rest during the backward pass. With checkpoints every √n layers, activation memory falls to O(√n) for roughly one extra forward pass of compute. In practice people checkpoint each transformer block. Korthikanti et al. showed selective recomputation, of only the cheap-to-recompute but memory-hungry attention parts, keeps most of the saving with far less overhead.

3. The practice project

(a) The laptop and free-GPU version

Two parts. First, build the primitives yourself on CPU. PyTorch's gloo backend runs collectives between processes on one machine, so torchrun --nproc_per_node=4 gives you four "GPUs" on a laptop. Write a column-then-row parallel MLP, check it matches the unsharded result, then time all-reduce across message sizes:

import time, torch, torch.distributed as dist
dist.init_process_group("gloo")
rk, ws = dist.get_rank(), dist.get_world_size()
torch.manual_seed(0)                      # same full weights on every rank
d, h = 1024, 4096
W1, W2 = torch.randn(d, h) / d**.5, torch.randn(h, d) / h**.5
x = torch.randn(8, d)
s = h // ws
y = torch.relu(x @ W1[:, rk*s:(rk+1)*s]) @ W2[rk*s:(rk+1)*s]  # partial sum
dist.all_reduce(y)                        # the one forward all-reduce
if rk == 0:
    print("max err", (y - torch.relu(x @ W1) @ W2).abs().max().item())

for n in [2**k for k in range(10, 25, 2)]:
    t = torch.ones(n)
    dist.barrier(); t0 = time.perf_counter()
    for _ in range(5):
        dist.all_reduce(t)
    dt = (time.perf_counter() - t0) / 5
    if rk == 0:
        bw = 2 * (ws - 1) / ws * n * 4 / dt / 1e9
        print(f"{n*4/1e6:8.2f} MB {dt*1e3:8.2f} ms  busbw {bw:.2f} GB/s")
dist.destroy_process_group()

You should see a tiny max error (float rounding), and bus bandwidth rising with message size until it levels off. That curve is why frameworks bucket small gradients together.

Second, get two real GPUs. Kaggle notebooks offer a free two-GPU T4 option. Fine-tune a model of about 0.5–1B parameters with DDP, then FSDP, then FSDP plus activation checkpointing. For each, record peak memory per GPU (torch.cuda.max_memory_allocated), tokens per second, and the share of step time in NCCL kernels from torch.profiler. T4s talk over PCIe, not NVLink, so communication shows up clearly.

(b) The full version from the post

Shard a 7B model across 4 GPUs, profile NCCL, and tune the tensor-parallel layout.

• Hardware: one node with 4 × 80 GB GPUs (A100 or H100), ideally NVLink-connected. Rental is a few dollars per GPU-hour, so a day of experiments costs in the low hundreds of dollars.
• Configurations to compare: FSDP (full sharding) alone; TP = 4; TP = 2 × FSDP = 2; each with and without activation checkpointing and sequence parallelism. torchtitan, Megatron-Core or DeepSpeed all support these.
• What to profile: step time, tokens per second per GPU, model FLOPs utilisation, peak memory, and time in each collective type. Use the PyTorch profiler trace or Nsight Systems, and run nccl-tests first to know your link's real bandwidth.
• "Optimising TP boundaries" means making sure each block has only its column-then-row pair of all-reduces, that no extra gathers sneak in (for example around layer norms or the loss), and that communication overlaps with compute. At 7B on one NVLink node, expect FSDP alone to be competitive. TP earns its keep at larger scale, longer sequences or slower links. Measure rather than assume.

(c) Pitfalls

• Measuring the first steps. Warm-up, compilation and allocator growth distort them. Skip the first 10–20 steps.
• Silent fallback to slow transport. Check NCCL_DEBUG=INFO output for the transport actually used.
• Hangs. One rank skipping a collective (a data-dependent branch, an uneven last batch) deadlocks everyone. Set timeouts and drop the last uneven batch.
• Different random seeds for dropout or initialisation across TP ranks break correctness silently. Compare against a single-GPU run on a small config first.
• Checkpoint formats. Sharded checkpoints must be saved and reloaded with the same tooling, or converted. Test a resume early.
• Effective batch size changes with DP degree. Keep it fixed when comparing layouts, or the loss curves aren't comparable.

4. Reality check

5. In this course

• Module 06 · Day 24, GPipe: pipeline parallelism, micro-batches and the bubble, with a demo.
• Module 05 · Serving a trillion-parameter model: the inference side of sharding, and MoE routing, which needs expert parallelism.
• Module 06 · Day 25, Scaling laws: why anyone needs this much compute.
• Module 02 · TensorTonic: the CUDA and Triton tracks, for the kernels that sit under all of this.

6. Why it matters for robotics

7. Self-check

8. Courses & content

Start with the Ultra-Scale Playbook: it covers every technique on this page in one place, with real measurements, and uses PyTorch like the project does.

Full course map for all 12 stages →

9. Papers & docs

• Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism (Shoeybi et al., 2019)
• ZeRO: Memory Optimizations Toward Training Trillion Parameter Models (Rajbhandari et al., 2020)
• Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM (Narayanan et al., 2021)
• Reducing Activation Recomputation in Large Transformer Models (Korthikanti et al., 2022)
• Training Deep Nets with Sublinear Memory Cost (Chen et al., 2016)
• PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel (Zhao et al., 2023)
• PyTorch FSDP tutorial and DeepSpeed ZeRO tutorial