Distributed Training: FSDP, DeepSpeed, and Megatron

Here's a number that bothered me the first time I sat down with it. A 70-billion-parameter model, trained the ordinary way with mixed-precision Adam, needs about 1.1 terabytes of GPU memory just for its weights, gradients, and optimizer states. An H100 has 80 gigabytes. You can arrange these GPUs in whatever topology you want; they still each have 80 GB, and no amount of cleverness changes that the model doesn't fit in one.

So how does anyone train these things? The answer is the subject of this post, and it's richer than "buy more GPUs." You can duplicate the model across machines (data parallelism), but that doesn't help if even one copy doesn't fit. You can slice the model into pieces and scatter the pieces across the cluster (model parallelism), but now every forward pass has to ferry data across the network. Most real training runs do both, and increasingly a third thing on top.

The good news is that the underlying ideas are simpler than the acronym soup (FSDP, DDP, ZeRO-3, TP, PP, 3D parallelism) makes them sound. Once you see the per-GPU memory budget and where each byte is spent, the whole landscape collapses to a small number of choices about what to replicate, what to shard, and when to communicate.

Prerequisites

This post assumes you've read:

• Gradients and Learning for how forward, loss, backward, and optimizer step fit together.
• GPUs and Floating-Point Formats so "fp16" and "fp32 master weights" mean something.
• Training View vs. Inference View so you know training is compute-bound and have a mental model of a forward pass over a full sequence.

If "what lives in GPU memory during training" feels fuzzy, the first few sections here will fix that before we get to the parallelism strategies.

Why one GPU isn't enough

Let's count the bytes. For a model with $P$ parameters trained in mixed precision with Adam (the default for every major open LLM), the GPU has to hold:

• Weights in fp16: $2P$ bytes.
• Gradients in fp16: $2P$ bytes.
• Optimizer states: an fp32 master copy of the weights ($4P$), first moment ($4P$), and second moment ($4P$). So $12P$ bytes.
• Activations for the backward pass. This one varies with batch size, sequence length, and architecture, but it's often on the same order as the weights.

That's roughly $16P$ bytes of model state plus activations. For 70B parameters: $16 \times 70 \times 10^9 = 1.12$ TB just for the model state. Even with aggressive activation checkpointing, you're at a terabyte-class problem.

An H100 has 80 GB. An H100 can hold about a 5B model trained the normal way. Training a Llama 70B or Mixtral or Qwen-max on a single GPU isn't an engineering challenge; it's a physical impossibility.

You could put different layers on different cards. That works, and we'll see how below. But the more interesting insight is that the memory is being spent on very non-uniform things. The optimizer states alone are $12P$ bytes, which is six times the weights. If we can get those off most of the GPUs, we've already cut 75% of the pressure. That observation is the entire ZeRO paper.

Pulling that thread gives you a family of techniques that shard different pieces of the model state across GPUs, each trading communication for memory at a different ratio. The rest of this post is a tour of those techniques.

Data parallelism: the cheap thing that works until it doesn't

The simplest distributed training you can do: put a complete copy of the model on every GPU, give each GPU a different chunk of the batch, and have them synchronize gradients at the end of each step. This is data parallelism, and in PyTorch it's DistributedDataParallel (DDP).

The mechanics are clean. Each of $N$ GPUs has an identical replica of the model weights. At each step, the global batch gets split into $N$ micro-batches, one per GPU. Forward and backward run independently on each GPU, and at the end each GPU has its own gradient tensor. Then all $N$ gradient tensors get averaged across the cluster via an all-reduce operation, so every GPU ends up with the same averaged gradient. Each GPU runs the optimizer step locally. Because they started from the same weights and just applied the same averaged gradient, they're bit-identical after the step.

DDP scales throughput linearly (up to network bandwidth) because all the GPUs are doing actual compute in parallel. The only new operation is that one all-reduce per step, which modern interconnects (NVLink, InfiniBand) can pipeline with the backward pass so it's mostly free.

The catch: each GPU still holds the entire model. Weights, gradients, optimizer states, all of it. DDP distributes the data. It does not distribute the model. If the model doesn't fit on one GPU in the first place, DDP does nothing for you.

For 7B models on 8 H100s, DDP is still the right answer. For anything bigger, you need to shard the model itself.

The ZeRO insight: you don't need a full copy on every GPU

In 2019, Microsoft's DeepSpeed team made a deceptively simple observation. During a DDP step, every GPU holds identical weights, identical gradients, and identical optimizer states. All of that duplication is pure waste. If you shard the state across $N$ GPUs, you divide each GPU's memory footprint by $N$.

The wrinkle is timing. You can shard weights, but at some point each GPU needs the full weight matrix to do a forward pass. You can shard gradients, but the optimizer step needs the gradient for each parameter on the GPU that owns that parameter. ZeRO figures out the exact communication pattern that makes each stage work.

The technique comes in three flavors, each cutting more memory at the cost of a little more communication.

• Stage 1 shards the optimizer states. They're the biggest chunk ($12P$ out of $16P$) and the only one that's needed just once per step (during the optimizer step), so sharding them is cheap: the cost is a reduce-scatter at the end of the backward pass so each GPU gets only its shard of the gradient. No new communication during forward or backward.
• Stage 2 shards gradients as well. Now each GPU only holds gradients for the parameters it owns. The backward pass writes full gradients temporarily, then a reduce-scatter drops them down to per-rank shards. Memory savings: about 2x over Stage 1. Cost: the reduce-scatter now dominates a bigger fraction of the gradient state.
• Stage 3 shards the weights themselves. This is the big one. Each GPU only permanently stores $P/N$ parameters. Before a layer can do its forward or backward pass, its full weights have to be gathered from all ranks (an all-gather). After the compute, the weights can be thrown away and regathered next time. Memory drops to roughly $16P / N$ per GPU.

The memory wins stack: Stage 3 on 64 GPUs is a ~64x reduction in per-GPU model state versus DDP. The communication volume grows, but for reasonable batch sizes the math works out: more parameters that fit means bigger models that train.

One more beat. PyTorch's FSDP (Fully Sharded Data Parallel) is essentially ZeRO-3 rewritten as a native PyTorch feature. It wraps a module, shards its parameters across ranks, and gathers them just-in-time during the forward pass. DeepSpeed predates FSDP and still has a few features FSDP lacks, but the two tools are converging fast, and most new LLM training recipes written in pure PyTorch reach for FSDP by default.

FSDP's rhythm: all-gather, compute, reshard

To feel what FSDP is actually doing, picture a 4-layer transformer on 4 GPUs. Each GPU holds 1/4 of every layer's weights in permanent storage. When the forward pass starts, layer 1's weights are reassembled via all-gather: every GPU temporarily holds the full layer-1 weight tensor. The matmul runs. Then the gathered weights get dropped, leaving only the sharded 1/4 behind. Next, layer 2 goes through the same gather-compute-reshard cycle. And so on.

The backward pass does the same thing in reverse, and along the way it accumulates gradients in full and then reduce-scatters them down so each rank keeps only its own shard of the gradient tensor. The result is that peak memory scales as $16P/N$ plus one layer's worth of transient gathered parameters, which is a tiny fraction of what DDP needs.

The tradeoff is communication volume. In DDP you do one all-reduce per step. In FSDP you do one all-gather per layer in the forward pass, one all-gather per layer in the backward, and one reduce-scatter per layer. That's on the order of $3L$ collective operations per step where $L$ is the number of layers. The communication volume in bytes is actually the same as DDP's all-reduce (you're moving the same parameters around), it's just spread across many small ops instead of one big one. Modern FSDP implementations overlap these collectives with the compute of neighboring layers, so the wall-clock penalty can be almost zero if your interconnect is fast enough.

A useful way to hold this in your head: FSDP is the memory-optimal version of data parallelism. You still do data parallelism semantically (each rank processes a different micro-batch), but you pay no memory cost for replicating the model, because you don't replicate it. The model exists, in any complete form, only for a few milliseconds at a time during compute.

Minimal math: where the memory actually goes

Before the worked example, let me put the memory numbers in one place. For a model with $P$ parameters, trained in mixed precision with Adam, on $N$ GPUs:

Activations are on top of this and depend on batch size, sequence length, and how much you checkpoint. For a 70B model on 64 H100s with FSDP, the model state alone drops to about 17 GB per GPU, which is the difference between "doesn't fit" and "fits with room for activations and a reasonable batch."

Play with the controls. A 70B model on 8 GPUs with DDP doesn't fit. On 64 GPUs with FSDP it does, comfortably. A 405B model needs ~128 GPUs just to fit the model state under FSDP alone (at $16P/N$, that's ~50 GB per GPU), and even then there's barely room for activations. That's the cue that you need tensor parallelism on top, which we'll get to.

Worked example

Here's the kind of napkin math someone running this training does.

The model. 70 billion parameters. Mixed precision with Adam. That's $16 \times 70 \times 10^9 = 1.12$ TB of model state, plus activations.

The hardware. 64 H100s, each with 80 GB of HBM3. Interconnect: 8 GPUs per node with NVLink (900 GB/s intra-node), nodes connected by InfiniBand (roughly 100 GB/s). Total cluster memory: $64 \times 80 = 5.12$ TB. That's enough for the model plus plenty of activations, but only if we shard.

Strategy. FSDP across all 64 GPUs.

Per-GPU budget under FSDP.

• Weights: $2 \times 70 \times 10^9 / 64 \approx 2.2$ GB
• Gradients: $2 \times 70 \times 10^9 / 64 \approx 2.2$ GB
• Optimizer states: $12 \times 70 \times 10^9 / 64 \approx 13.1$ GB
• Model state subtotal: ~17.5 GB

That leaves about 62 GB per GPU for activations, transient gathered layers, buffers, and a micro-batch. Plenty of room. With a sequence length of 8192 and activation checkpointing, each GPU can handle a micro-batch of a few dozen sequences, which combined across 64 GPUs gives a global batch of a couple of thousand. Comfortably in the zone where convergence is well-studied.

Throughput. A well-tuned FSDP run on this cluster can hit something like 40-60% of peak H100 fp16 TFLOPS (the rest is lost to activations, communication overhead, and pipeline bubbles). At 400 TFLOPS effective per GPU, across 64 GPUs, you're at ~25 PFLOPS sustained. Training 70B on 1.4T tokens (Chinchilla-optimal) would take a few weeks at this scale. In practice, Meta trained Llama 3 70B on ~15T tokens, roughly 10× the Chinchilla-optimal token count, which took correspondingly longer but used the same FSDP mechanics on a much bigger cluster.

The broader point: once you've decided on FSDP, the remaining decisions (global batch size, micro-batch per GPU, gradient accumulation, learning rate schedule) are ordinary training hyperparameters. The hard part was getting the model to fit. That's what ZeRO-3/FSDP bought us.

When FSDP isn't enough: Megatron and tensor parallelism

ZeRO-3 gets you very far. If you have enough GPUs, it gets you to 405B. But there are two situations where even FSDP hits a wall.

One: a single layer is too big. Modern LLMs have FFN blocks with $d_{ff} \approx 4 d_{model}$ where $d_{model}$ can be 16,384 or more. That makes one FFN weight matrix a $16{,}384 \times 65{,}536$ tensor, about 2 GB of parameters for that single layer. The all-gather to materialize that layer's weights crosses the interconnect, and if you're pushing 2 GB over a 100 GB/s link for every layer on every forward pass, the link becomes the bottleneck.

Two: the micro-batch is too small. FSDP shards the model but not the compute. When you use many GPUs, each one processes $1/N$ of the batch. Below some threshold, the GPU is just idle and arithmetic intensity tanks.

The fix for both is to shard the computation itself, not just the state. That's what the Megatron-LM team built in 2019, and it comes in two flavors: tensor parallelism and pipeline parallelism.

Tensor parallelism (TP) takes a single matmul and slices its weight matrix along the output dimension. Each GPU computes a slice of the output. A subsequent all-gather (or all-reduce, depending on how the next layer is split) combines the slices. This works because matrix multiplication distributes naturally: $(X \cdot W)$ split column-wise gives you pieces of the output that just need to be concatenated.

Tensor parallelism is fast in practice because the communication happens inside a single node, over NVLink. You typically set TP=8 on a DGX box (one TP group per node) and call it done, so you never have to cross the slow InfiniBand network to compute a single layer.

Pipeline parallelism (PP) splits the layers of the model across different GPU groups. Stage 1 holds layers 1-8, stage 2 holds layers 9-16, and so on. A forward pass walks through the stages: micro-batch 1 goes through stage 1, then stage 2, while stage 1 picks up micro-batch 2. This is a wavefront pattern. The pipeline fills, runs steady-state, then drains.

Pipeline parallelism has a classic inefficiency called the bubble. At startup and shutdown the pipeline isn't full, and stages sit idle. Good schedules (1F1B, interleaved 1F1B) minimize bubble size by keeping more micro-batches in flight. At Llama-3 scale, PP is the last resort, used only when TP and data parallelism aren't enough on their own.

3D parallelism: how 16K GPUs train one model

At frontier scale, no single parallelism strategy is enough. You compose them into a 3D grid.

• Data parallelism (or ZeRO/FSDP within data parallelism) across one axis. Each data-parallel replica sees a different micro-batch.
• Tensor parallelism within each replica, usually across the 8 GPUs in a node, using NVLink.
• Pipeline parallelism across nodes, with each stage holding a contiguous range of layers.

Every GPU in the cluster is identified by a 3-tuple (dp_rank, tp_rank, pp_stage). For Llama 3's 405B run, Meta reported something like $N_{DP} \times N_{TP} \times N_{PP}$ arranged to balance all three.

The three axes solve genuinely different problems. DP scales throughput and distributes the optimizer state. TP keeps a single enormous matmul from exceeding one GPU's memory or compute. PP lets you put tens or hundreds of layers into a training run without each stage having to hold the whole model.

Getting the ratios right is where real expertise lives. Too much PP and the bubble eats your throughput. Too much TP and the cross-GPU communication dominates. Too much DP and you need larger global batches than converge well. The published recipes from Meta, Mistral, Alibaba, and DeepSeek all land in a similar zone, with TP around 8, PP between 8 and 32, and DP making up the rest to fill the cluster.

Implementation sketch

For the code block, the smallest thing that actually teaches FSDP is one wrapper around a model and one training step. Everything else (dataset, tokenizer, eval) is orthogonal.

import torch
import torch.distributed as dist
from torch.distributed.fsdp import (
    FullyShardedDataParallel as FSDP,
    MixedPrecision,
    ShardingStrategy,
)
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
import functools
 
# --- Init process group: one process per GPU ---
dist.init_process_group(backend="nccl")
rank = dist.get_rank()
torch.cuda.set_device(rank)
 
# --- Build the raw model on GPU 0, CPU elsewhere ---
# (In practice you'd use meta-tensors to avoid ever materializing the full model)
from my_transformer import TransformerDecoder, TransformerBlock
model = TransformerDecoder(vocab=50257, d_model=4096, n_layers=32, n_heads=32)
 
# --- Wrap with FSDP ---
# auto_wrap_policy tells FSDP to shard each TransformerBlock separately,
# so the all-gather scope is one block at a time, not the whole model.
auto_wrap_policy = functools.partial(
    transformer_auto_wrap_policy,
    transformer_layer_cls={TransformerBlock},
)
 
mp_policy = MixedPrecision(
    param_dtype=torch.bfloat16,
    reduce_dtype=torch.bfloat16,
    buffer_dtype=torch.bfloat16,
)
 
model = FSDP(
    model,
    auto_wrap_policy=auto_wrap_policy,
    sharding_strategy=ShardingStrategy.FULL_SHARD,  # == ZeRO-3
    mixed_precision=mp_policy,
    device_id=rank,
)
 
# --- Optimizer: sees only the local shards, not the full params ---
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
 
# --- One training step ---
def train_step(batch):
    # FSDP auto-gathers each block's params at forward time,
    # then reshards as soon as compute is done.
    logits = model(batch["input_ids"])
    loss = torch.nn.functional.cross_entropy(
        logits.reshape(-1, logits.size(-1)),
        batch["labels"].reshape(-1),
    )
    # Backward pass: gradients get reduce-scattered across ranks,
    # so each rank ends up with only its shard of the gradient.
    loss.backward()
    optimizer.step()       # updates only the local param shard
    optimizer.zero_grad()
    return loss.item()

Notice how little of this is "distributed code." The FSDP(model, ...) wrapper is the whole thing. Everything else (forward, backward, optimizer step) reads exactly like single-GPU training. PyTorch hides all the all-gathers and reduce-scatters behind hooks on the module's forward and backward methods. That's a large part of why FSDP displaced DeepSpeed for a lot of teams: the code looks like the code you already wrote.

The one conceptual piece that's not obvious from the snippet is the auto_wrap_policy. FSDP shards one "unit" at a time. If the unit is the whole model, you get a single gigantic all-gather at every forward pass and no overlap with compute. If the unit is one transformer block, you get many smaller gathers that can be overlapped with the compute of neighboring blocks. Picking the right granularity is the main knob you turn when tuning FSDP for throughput.

Misconceptions

"FSDP is model parallelism." Not quite. FSDP is still data parallelism in the sense that every rank runs the full computation on a different micro-batch; every GPU does a forward pass on every layer. What FSDP shards is the storage of the parameters, not the compute. Tensor parallelism is model parallelism proper: each GPU computes only a slice of each operation, and the results have to be combined. They compose, but they aren't the same thing.

"Adding more GPUs always helps." Only up to a point. Data parallelism (including FSDP) scales until the global batch size is so large that gradients stop being informative. There's a sweet spot around a few million tokens per batch for current LLMs, above which you're doing extra compute for less learning per step. And communication overhead grows with $N$: even if the collective ops are latency-hidden, at some point the network becomes the bottleneck. Frontier runs on 16K+ GPUs spend serious engineering effort on bandwidth routing, topology-aware scheduling, and failure handling because adding GPUs naively stops helping around a few thousand.

"ZeRO-3 and FSDP are different things." Same idea, different implementations. ZeRO-3 is the DeepSpeed name for "shard parameters, gradients, and optimizer states across ranks." FSDP is PyTorch's in-tree implementation of the same algorithm. There are API differences and a few feature differences (DeepSpeed has better inference support, FSDP has tighter PyTorch integration, FSDP2 further polishes the API), but conceptually they're the same technique. If someone says "we use ZeRO-3" and someone else says "we use FSDP," they mean the same strategy.

What's next

The next post covers supervised fine-tuning, how you take a pretrained base model and teach it to follow instructions. It's the first step in the post-training pipeline that turns a next-token predictor into something you'd actually want to talk to.

Additional reading (and watching)

• Korthikanti, V., et al. (2022). Reducing Activation Recomputation in Large Transformer Models. arXiv:2205.05198. Good source for the per-GPU memory accounting, including activations under selective checkpointing.

• Rajbhandari, S., et al. (2019). ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. arXiv:1910.02054. The foundational paper that established the three stages and the per-GPU memory math this whole post is built on.

• Li, S., et al. (2020). PyTorch Distributed: Experiences on Accelerating Data Parallel Training. VLDB 2020. The DDP paper — documents the gradient bucketing and overlap tricks that make naive data parallelism fast in PyTorch.

• NVIDIA. (2020). NCCL: Optimized primitives for collective multi-GPU communication. Documentation. The library behind every all-reduce, all-gather, and reduce-scatter in modern training stacks.

• Zhao, Y., et al. (2023). PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel. VLDB 2023. The PyTorch team's FSDP paper, covering the design choices, auto-wrap policy, and how FSDP composes with other parallelism axes.

• Touvron, H., et al. (2023). LLaMA 2: Open Foundation and Fine-Tuned Chat Models. arXiv:2307.09288. Training recipe for Llama 2, which used FSDP-style sharding at 70B scale.

• Meta AI. (2024). The Llama 3 Herd of Models. arXiv:2407.21783. The engineering appendix of this paper is the best single source for what frontier-scale distributed training actually looks like in practice, including the 3D parallelism recipe, failure rates, and bandwidth accounting.

• Shoeybi, M., et al. (2019). Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. arXiv:1909.08053. The paper that established tensor parallelism as a practical technique; the Narayanan et al. (2021) follow-up Efficient Large-Scale Language Model Training on GPU Clusters adds pipeline parallelism and the 3D recipe.

• Huang, Y., et al. (2019). GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism. NeurIPS 2019. The original pipeline-parallelism paper, with the classic bubble-time analysis.