Distributed Training: FSDP, DeepSpeed, and Pipeline Parallelism

Suppose a lab at IIT Bombay is fine-tuning a 7-billion-parameter multilingual Indic model — something in the spirit of AI4Bharat's Airavata — on a single DGX node with eight A100 GPUs, 80 GB each, 640 GB pooled across the node. The multi-GPU chapter you have already read showed how eight replicas of a model synchronize their gradients with an all-reduce so they stay identical after every step. That machinery assumes each GPU can hold one full copy of the model, its gradients, and its optimizer state. At 7B parameters, that assumption breaks — not because the node lacks memory, but because no single GPU in it does. This chapter is about what happens once replication itself becomes the bottleneck, and the three techniques built to route around it: DeepSpeed's ZeRO optimizer, PyTorch's FSDP, and pipeline parallelism.

Where the 112 GB comes from

Start from first principles. Training with mixed precision and the Adam optimizer, each parameter carries not one but five separate numbers in memory, because Adam needs its own running statistics and because mixed precision keeps a full-precision master copy alongside the half-precision working copy used for the matrix multiplies. Rajbhandari, Rasley, Ruwase, and He formalized this accounting in their ZeRO paper (Microsoft, SC20, 2020), using Ψ for the parameter count:

• fp16 parameters: 2 bytes/param — used for the forward and backward matmuls
• fp16 gradients: 2 bytes/param — accumulated during backward
• fp32 master copy of parameters: 4 bytes/param — the numerically stable copy the optimizer actually updates
• fp32 Adam momentum (first moment): 4 bytes/param
• fp32 Adam variance (second moment): 4 bytes/param

That totals 16Ψ bytes. For Ψ = 7×10⁹, that is 112×10⁹ bytes ≈ 112 GB — per replica, before a single activation tensor is stored. Plain data parallelism, of the kind the multi-GPU chapter covers, replicates all 112 GB onto every one of the eight GPUs. It does not fit on one 80 GB card, so DDP alone is a dead end here regardless of how many GPUs the node has, because DDP never divides that 112 GB — it only divides the batch.

DeepSpeed ZeRO: shard the state, not the batch

ZeRO (Zero Redundancy Optimizer) keeps the DDP communication pattern's spirit — every GPU still processes a different slice of the batch — but stops replicating the 16Ψ bytes wholesale. It partitions that state across the N data-parallel ranks in three increasingly aggressive stages, each GPU permanently owning only 1/N of a given category and fetching the rest just-in-time from its owning rank.

Stage 1 (Pos) shards only the 12Ψ bytes of optimizer state (fp32 master + momentum + variance), keeping fp16 params and fp16 grads (4Ψ) replicated. Per-GPU memory = 4Ψ + 12Ψ/N.

Stage 2 (Pos+g) additionally shards the fp16 gradients, so only the fp16 parameters (2Ψ) stay replicated. Per-GPU memory = 2Ψ + 14Ψ/N.

Stage 3 (Pos+g+p) shards the parameters too — nothing is permanently replicated. Per-GPU memory = 16Ψ/N.

Plugging in the IIT Bombay scenario, Ψ = 7×10⁹, N = 8:

Stage 1: 4(7) + 12(7)/8  =  28.0 + 10.5  =  38.5 GB/GPU
Stage 2: 2(7) + 14(7)/8  =  14.0 + 12.25 =  26.25 GB/GPU
Stage 3: 16(7)/8         =  14.0 GB/GPU

Stage 3 takes the model from "impossible on 640 GB pooled across 8 cards" to "14 GB per card, with 66 GB left over on each A100 for activations and KV-style intermediate tensors." This is not free — fetching sharded state back on demand costs bytes on the wire, which is the next thing to get exactly right rather than wave at. In plain DDP, a ring all-reduce of the gradients moves about 2Ψ bytes off each GPU. ZeRO-1 and ZeRO-2 preserve that same ~2Ψ figure: gradients are reduce-scattered (Ψ) and updated parameters are all-gathered back (Ψ), which is why the ZeRO paper's headline result is that stage 2 matches DDP's communication volume while cutting memory roughly 8×. Stage 3 adds one more all-gather, because parameters themselves are no longer resident — a GPU must pull in the full parameter tensor before it can run the forward pass through a given layer, and pull it in again before backward. That pushes total traffic to about 3Ψ, a 1.5× increase over the DDP baseline. Sharding state to fit memory has a real, quantified communication price, and that price is exactly why stage 3 is the option of last resort rather than the default.

A minimal DeepSpeed setup for this expresses the stage as configuration, not code:

import deepspeed
import torch
import torch.nn as nn

ds_config = {
    "train_batch_size": 256,
    "fp16": {"enabled": True},
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {"device": "cpu"},
        "overlap_comm": True,
        "contiguous_gradients": True
    }
}

model = nn.Linear(4096, 4096)                            # stand-in for the real 7B model
data_loader = [torch.randn(8, 4096) for _ in range(10)]  # stand-in batches

model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    model_parameters=model.parameters(),
    config=ds_config
)

for step, batch in enumerate(data_loader):
    loss = model_engine(batch)
    model_engine.backward(loss)
    model_engine.step()

The offload_optimizer line is worth noting even without deriving its numbers: ZeRO-Infinity (Rajbhandari et al., 2021) extends the same partitioning idea to push shards onto CPU RAM or even NVMe, trading PCIe/NVMe bandwidth for the ability to train models that would not fit in GPU memory even fully sharded across the node.

FSDP: the same idea, a different unit of sharding

PyTorch's Fully Sharded Data Parallel, described by Zhao et al. in "PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel" (VLDB 2023), implements the same memory arithmetic as ZeRO stage 3 — 16Ψ/N per GPU, ~3Ψ communication — but organizes the sharding differently, and that difference is worth being able to name precisely rather than treating FSDP and ZeRO-3 as interchangeable synonyms.

DeepSpeed shards at the granularity of individual parameter tensors, tracked by a persistence/prefetch engine that decides which shards to fetch ahead of when they are needed. FSDP instead wraps the model in a tree of "FSDP units" (by default, roughly one per top-level submodule, controlled by an auto_wrap_policy). Every parameter inside a unit is flattened and concatenated into a single 1-D tensor called a FlatParameter, and that FlatParameter — not the individual weight matrices inside it — is what gets sharded 1/N per GPU. Just before an FSDP unit's forward (and again before its backward) runs, every GPU issues an all-gather to reconstruct that unit's full FlatParameter locally, runs the computation, and then immediately frees (re-shards) it. Only one unit's worth of full parameters needs to be materialized in memory at a time, which is what keeps peak memory close to the 16Ψ/N figure despite briefly holding un-sharded weights during compute.

import functools
import torch.nn as nn
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy

# (assumes torch.distributed.init_process_group has already run,
#  as in the multi-GPU/multi-node chapter)

class TransformerBlock(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.ln = nn.LayerNorm(dim)
        self.attn = nn.MultiheadAttention(dim, num_heads=8, batch_first=True)
        self.mlp = nn.Sequential(nn.Linear(dim, 4 * dim), nn.GELU(), nn.Linear(4 * dim, dim))

    def forward(self, x):
        h = self.ln(x)
        x = x + self.attn(h, h, h)[0]
        return x + self.mlp(x)

class DemoTransformer(nn.Module):
    def __init__(self, dim=4096, depth=32):
        super().__init__()
        self.blocks = nn.ModuleList([TransformerBlock(dim) for _ in range(depth)])

    def forward(self, x):
        for block in self.blocks:
            x = block(x)
        return x

model = DemoTransformer().cuda()
policy = functools.partial(size_based_auto_wrap_policy, min_num_params=20_000_000)
fsdp_model = FSDP(model, auto_wrap_policy=policy)

Here size_based_auto_wrap_policy tells FSDP to close off a new unit once the parameters it has accumulated exceed 20 million — in practice this tends to land close to one transformer block per unit, which is exactly the granularity you want: small enough that only a modest slice of the 112 GB is ever fully materialized at once, large enough that the all-gather isn't so fine-grained it can't overlap with compute.

Pipeline parallelism: partitioning the computation itself

ZeRO and FSDP both leave every GPU doing the exact same FLOPs as an unsharded replica would — they only change where the bytes live between compute steps. Pipeline parallelism is a genuinely different move: it cuts the model's layers into p contiguous stages, assigns each stage to one GPU, and each GPU only ever runs the forward and backward math for its own slice of layers. Activations flow stage-to-stage like a network packet, not like a gradient.

The naive version of this idea is wasteful: send one microbatch through stage 0, then stage 1, then stage 2 — while stages 1 and 2 sit idle waiting, and stage 0 sits idle once it has handed off. GPipe (Huang et al., NeurIPS 2019) fixes this by splitting each training batch into m microbatches and streaming them through the pipeline so multiple stages are busy on different microbatches simultaneously — but a start-up "fill" and a shutdown "drain" are unavoidable, and that unavoidable idle time is the pipeline bubble.

Derive the bubble fraction directly. Let each stage take time t to run one microbatch's forward pass and the same t for backward (a standard simplifying assumption). Stage i cannot start its forward on microbatch j until stage i−1 has finished handing it off, so stage i's forward of microbatch j begins at time i·t + (j−1)·t; the last stage (index p−1) finishes its forward of the last microbatch (j = m) at time (p−1)·t + (m−1)·t + t = (m+p−1)·t. Backward runs as a mirror image, stage p−1 first, flowing back to stage 0, taking another (m+p−1)·t. Total wall-clock time is therefore:

T_total = (m + p - 1) * (t_f + t_b)          [with t_f = t_b = t: T_total = 2(m+p-1)t]

Ideal (bubble-free) work per stage = m * (t_f + t_b)
Bubble time per stage              = (p - 1) * (t_f + t_b)
Bubble fraction of wall clock      = (p - 1) / (m + p - 1)

For p = 4 stages and m = 8 microbatches: bubble fraction = 3/11 ≈ 27.3%. Increase to m = 32 microbatches with the same 4 stages: 3/35 ≈ 8.6%. Finer microbatching amortizes the fixed fill/drain cost over more useful work — which is exactly why GPipe's paper recommends m ≫ p, and why the commonly quoted shortcut formula (p−1)/m (here 3/8 = 37.5% and 3/32 = 9.4%) is only a good approximation once m is large relative to p; at m = 8 it overstates the true bubble by roughly 37-38%.

Microbatching more finely is not free either: GPipe must hold activations for every microbatch still in flight, so its original design pairs pipelining with activation checkpointing — recomputing each stage's activations during backward instead of storing them — to keep peak activation memory bounded by the number of layers per stage rather than by m. PipeDream (Narayanan et al., SOSP 2019) attacks the same problem from the scheduling side with a 1-forward-1-backward (1F1B) schedule: instead of running all m forward passes before any backward starts, a stage alternates one forward and one backward as soon as dependencies allow, which caps the number of unretired microbatches at roughly p rather than m and lowers peak activation memory without sacrificing the throughput gained from pipelining.

The misconception: sharding is not splitting the computation

The mistake students make almost every time they meet ZeRO-3 or FSDP for the first time is assuming it works like pipeline or tensor parallelism — that GPU 3 must be "handling different layers" the way it does in the pipeline diagram above. It is not. Under ZeRO-3/FSDP, every one of the 8 GPUs still runs the forward and backward pass through every one of the 32 transformer blocks in the model, for its own slice of the batch, exactly as in plain DDP. What differs from DDP is only where the 112 GB of parameter, gradient, and optimizer bytes physically sit between those computations: sharded across GPUs and reassembled just-in-time, rather than fully replicated. FSDP and ZeRO-3 are storage-partitioning techniques layered on top of data parallelism; the per-GPU FLOP count is identical to DDP's. Pipeline parallelism (and tensor parallelism, which slices individual weight matrices) are computation-partitioning techniques — a GPU running stage 1 of a pipeline genuinely never executes the matrix multiplies belonging to stage 3. Conflating the two leads to real engineering mistakes: someone who thinks FSDP reduces per-GPU compute will be surprised that an FSDP job on 8 GPUs doesn't run faster per step than a DDP job that fits in memory — FSDP's payoff is fitting a model that otherwise couldn't run at all, not accelerating one that already could.

Active recall

Attempt these before reading the answers.

1. A 13-billion-parameter model is trained in mixed precision with Adam across N = 16 GPUs. What is the per-GPU memory footprint under ZeRO stage 1, and under stage 3?
2. In the pipeline example above (p = 4, m = 8), suppose the cluster is expanded to p = 8 stages while m stays at 8 microbatches. Compute the new bubble fraction, and explain what else changes besides throughput.
3. Why does ZeRO stage 2 keep the same ~2Ψ communication volume as plain DDP, while stage 3 needs ~3Ψ?
4. A colleague says: "We switched our 8-GPU job from DDP to FSDP and step time barely changed — so FSDP isn't actually saving us anything." What is wrong with the inference?
5. In FSDP, why is a single FlatParameter, rather than each individual weight tensor, the unit that gets all-gathered and re-sharded?
6. If a team increases GPipe's microbatch count from m = 8 to m = 64 at p = 4, what happens to the bubble fraction, and what is the cost of doing so that the bubble-fraction formula alone doesn't show?

Answers.

1. Ψ = 13×10⁹. Stage 1: 4Ψ + 12Ψ/N = 4(13) + 12(13)/16 = 52 + 9.75 = 61.75 GB/GPU — still too large for an 80 GB card once activations are added. Stage 3: 16Ψ/N = 16(13)/16 = 13 GB/GPU, comfortably leaving room for activations and a larger batch.

2. Bubble fraction = (p−1)/(m+p−1) = 7/15 ≈ 46.7%, up sharply from 27.3%. This is the full ripple, not just the headline number: (a) each device now holds only 1/8 of the layers instead of 1/4, so per-GPU parameter and activation memory for the model itself drops; (b) the number of stage-to-stage handoffs (pipeline boundaries) is p−1, so it more than doubles in frequency (3 boundaries at p=4 versus 7 at p=8), even though the tensors crossing each boundary are the same size; (c) because almost half of every device's time is now idle, the effective useful throughput per GPU falls — adding pipeline stages without adding microbatches to match makes the bubble worse, not better, which is exactly why p and m are tuned together rather than independently.

3. Stage 2 still reduce-scatters the gradients (Ψ bytes) and all-gathers the updated fp16 parameters back (Ψ bytes) — 2Ψ total, the same shape as DDP's all-reduce, because parameters are still held in full on every GPU between steps. Stage 3 additionally has to all-gather the full parameter tensor for a unit before it can even run that unit's forward pass (and again before backward), since no GPU holds a complete copy at rest — that extra gather adds roughly another Ψ, for ~3Ψ total.

4. FSDP's benefit is memory, not speed — per-GPU FLOP count is identical to DDP by design, so unchanged step time on a job that already fit in memory under DDP is exactly the expected, correct outcome, not a sign FSDP failed. The payoff shows up only when the model is too large to fit under DDP at all, or when the freed memory is used to increase batch size, sequence length, or model size — none of which "step time on the same config" would capture.

5. All-gathering many small individual tensors means many small collective calls, which underutilizes interconnect bandwidth and adds per-call overhead; concatenating a whole unit's parameters into one contiguous FlatParameter lets FSDP issue one large all-gather (and one large reduce-scatter for gradients) per unit, which is far more bandwidth-efficient on NCCL-style collectives than many tiny ones would be.

6. Bubble fraction = (p−1)/(m+p−1) = 3/67 ≈ 4.5%, down from 27.3%. The cost the formula hides is memory: more microbatches in flight before the first backward starts means more activation tensors must be kept live simultaneously (or, if using checkpointing, more recomputation work during backward) — the bubble shrinks, but only by asking every stage to hold or regenerate a longer queue of unfinished microbatches, which is precisely the memory pressure that PipeDream's 1F1B schedule was designed to relieve.

Think About It

Think about this: How would you explain distributed training: fsdp, deepspeed, and pipeline parallelism to a friend who has never seen a computer? What real-world analogy would you use? Imagine you had to build a system using these concepts — what would be your first step? Try this: before moving on, write down three things you learned and one question you still have.

Key Takeaways — Summary and Recap

Let us recap what we covered: the core ideas behind distributed training: fsdp, deepspeed, and pipeline parallelism, how they connect to real-world applications, and why they matter for your journey in computer science. Remember these key points as you move forward. For competitive exam preparation (CBSE, JEE, BITSAT), focus on understanding the WHY behind each concept, not just the WHAT.