AI/ML Explained: Training Infrastructure

A 70B model won't fit on one GPU — just like a table too big for one machine. Your instinct is sharding. But training has three orthogonal axes to split along: split the data (like read replicas), split the model's layers (like a microservice pipeline), or split the big matrix inside a single layer (like splitting one huge JOIN across machines). Understanding distributed training is really about how these three axes combine.

The pain point is VRAM, not compute. A 70B model in fp16 is 140GB of parameters alone; add gradients and Adam optimizer state and each parameter costs ~16 bytes during training (the ZeRO card breaks this number down) — over 1TB total. An 80GB H100 can't hold it. The three parallelisms each solve one piece:

• Data Parallel (DP): every GPU holds a full copy of the model, only the training batch is split. Each GPU runs forward/backward independently, then all-reduce averages the gradients. Analogy: multiple read replicas each handling a batch, reconciling periodically. Requires the model itself to fit on one GPU;
• Pipeline Parallel (PP): split the model by layers into stages across GPUs, like an assembly line / CPU pipeline. The cost is the bubble — stage 1 finishes and waits for stage 2; you fill it by chopping the batch into micro-batches;
• Tensor Parallel (TP): split the big matrix multiply inside a single layer across GPUs (detailed in card 4). Communication-heaviest, done only over a machine's high-speed NVLink.

# Simplest data parallel: PyTorch DDP. Launch multi-process with torchrun
import torch, torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

dist.init_process_group("nccl")        # NCCL = NVIDIA's GPU collective comm library
rank = dist.get_rank()
torch.cuda.set_device(rank)

model = MyTransformer().to(rank)
model = DDP(model, device_ids=[rank])  # auto all-reduce of grads on backward

for batch in shard_loader(rank):     # each rank gets a different data slice
    loss = model(batch).loss
    loss.backward()                  # ← gradients averaged across GPUs here
    optimizer.step(); optimizer.zero_grad()
# Launch: torchrun --nproc_per_node=8 train.py

Pure data parallel has a huge waste: every GPU stores a full copy of the parameters, gradients, and optimizer state — N redundant copies, like "replicate the full dataset to every node" in distributed storage. ZeRO's insight: these states are idle most of the time, so why keep them full on every GPU? Shard them across GPUs so each holds only 1/N, and all-gather on demand when needed. It's fundamentally trading communication for VRAM — the same trade-off as using erasure coding instead of full replication.

First, unpack the "16 bytes per parameter" — that's the key to ZeRO. In mixed-precision training each parameter stores:

ZeRO comes in three stages, progressively cutting redundancy — more savings, more communication:

• ZeRO-1: shard the optimizer state (those 12 bytes). Cut the fattest part first, almost no extra communication — best bang for the buck;
• ZeRO-2: also shard the gradients. After backward, reduce-scatter gradients to the owning GPU, keep only your slice locally;
• ZeRO-3: shard the parameters themselves too. When forward/backward reaches a layer, all-gather the full params, and discard immediately after use. Per-GPU VRAM drops to about 16/N bytes/param — in theory, adding GPUs linearly lowers VRAM.

This is Microsoft DeepSpeed's ZeRO (Rajbhandari et al. 2019, arXiv 1910.02054). Its elegance: under the data-parallel programming model (each GPU still sees the illusion of a "full model"), it quietly eliminates the memory redundancy.

# DeepSpeed enables ZeRO via one JSON config; business code barely changes
ds_config = {
    "train_micro_batch_size_per_gpu": 4,
    "bf16": {"enabled": True},
    "zero_optimization": {
        "stage": 3,                 # shard optimizer + grads + params
        "offload_optimizer": {"device": "cpu"},  # can even offload to CPU RAM
        "stage3_prefetch_bucket_size": 5e7        # prefetch to hide comm latency
    }
}
import deepspeed
engine, _, _, _ = deepspeed.initialize(
    model=model, config=ds_config, model_parameters=model.parameters())
for batch in loader:
    loss = engine(batch).loss
    engine.backward(loss)   # ZeRO does reduce-scatter + all-gather here
    engine.step()

FSDP is PyTorch officially baking ZeRO-3 into the standard library. The mechanism is nearly equivalent; the difference is like third-party library vs built-in: DeepSpeed/ZeRO is bolted on (like requests), FSDP is co-designed with PyTorch's autograd, dispatcher, and memory allocator (like built-in urllib, but smoother). For most pure-PyTorch projects, FSDP is the default way to get ZeRO-3 without adopting a new ecosystem.

It solves the same pain as ZeRO-3 — eliminate data-parallel memory redundancy — but must blend seamlessly into native PyTorch. The mechanism in one line: shard at rest, gather in use, free after use.

• Wrap into units: split the model into FSDP units by layer/module. At rest, each unit's params keep only a 1/N shard on this GPU;
• Forward into a unit: all-gather to reconstruct that layer's full params → compute the layer → free immediately the non-local parts, keeping only 1/N;
• Backward likewise: all-gather once more to compute gradients, then reduce-scatter the gradient shards back to their owning GPUs;
• Comm-compute overlap: while computing layer i, prefetch layer i+1's params in the background, hiding all-gather latency behind compute — the key to running fast.

The authoritative reference is PyTorch FSDP (Zhao et al. 2023, arXiv 2304.11277), which details the co-design with PyTorch internals. auto_wrap_policy sets granularity: units too fine → many all-gathers, fragmented comm; too coarse → high peak VRAM per gather. Granularity itself is a continuous VRAM-vs-communication trade-off.

# PyTorch-native FSDP, no third-party library
import torch, functools
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy

# Auto-split units by Transformer block — granularity sets the comm/VRAM trade-off
policy = functools.partial(
    transformer_auto_wrap_policy,
    transformer_layer_cls={TransformerBlock})

model = FSDP(
    model,
    auto_wrap_policy=policy,
    device_id=torch.cuda.current_device())

# The training loop afterward looks almost like single-GPU — sharding is transparent
for batch in loader:
    loss = model(batch).loss
    loss.backward()      # FSDP handles all-gather / reduce-scatter here
    optimizer.step(); optimizer.zero_grad()

ZeRO and FSDP both assume a single layer is computable — they shard "who stores params," but reassemble the full layer on one GPU to compute. What if a single layer's matrix is itself too big for one GPU? Then you must split the single matrix multiply itself, letting multiple GPUs cooperate on one matmul. Analogy: one record too big for a single machine — you have to split that record's fields across machines (vertically splitting a single operator). This is tensor parallelism, and Megatron-LM is its classic implementation.

Megatron-LM (Shoeybi et al. 2019, arXiv 1909.08053) is clever about how to split the Transformer's two big operators so the split boundary lands outside the nonlinearity — so one forward needs only one sync. Take the MLP layer Y = GeLU(X·A)·B:

• Split the first matrix A by columns into [A₁, A₂]: X·A₁ and X·A₂ compute independently on two GPUs. Because GeLU is element-wise, GeLU(X·A₁) doesn't need A₂'s result — the nonlinearity completes locally on each GPU, no intermediate sync;
• Split the second matrix B by rows into [B₁; B₂]: each GPU computes a partial sum, then one all-reduce sums them into the full Y.

Why "column-split A, row-split B" and not the reverse? Because this combo keeps the only nonlinearity (GeLU) inside the split, so the whole layer's forward syncs just once at the end. The other split would require a sync before and after GeLU — doubling communication. Self-attention is even more natural: split by attention head, since each head computes independently — born for parallelism.

Tensor parallel communication is heavy (an all-reduce every layer), so it's done only over a single machine's NVLink (~900GB/s bandwidth), never across nodes. The final thousand-GPU training is 3D parallelism:

# The core idea of tensor parallel: column-split + row-split + all-reduce, in plain PyTorch
import torch, torch.distributed as dist
import torch.nn.functional as F

def tp_mlp(X, A_shard, B_shard):
    # A split by columns: this GPU holds only some columns of A → partial hidden dims
    H = F.gelu(X @ A_shard)        # GeLU is element-wise, local, no sync needed
    # B split by rows: this GPU holds the matching rows → a partial sum
    Y_partial = H @ B_shard
    # The one sync: add the partial sums across GPUs = full output
    dist.all_reduce(Y_partial, op=dist.ReduceOp.SUM)
    return Y_partial            # the whole layer's forward communicated once
# Production uses NVIDIA Megatron-LM / Megatron-Core, which packages this splitting

• ZeRO (Rajbhandari et al. 2019) — foundational analysis of memory redundancy and three-stage sharding
• Megatron-LM (Shoeybi et al. 2019) — how to split the Transformer for tensor parallelism, a must-read
• PyTorch FSDP (Zhao et al. 2023) — co-design of FSDP with PyTorch internals
• GPipe (Huang et al. 2018) — the original pipeline-parallelism paper with micro-batches and the bubble
• DeepSpeed ZeRO tutorial — config usage and engineering trade-offs per stage
• PyTorch FSDP tutorial — hands-on auto_wrap_policy, state_dict, and more