DeepSpeed ZeRO: Extreme Memory Efficiency

Picture a small Indic-language modeling team, four A10G GPUs on a rented cluster, each card carrying 24 GB, trying to pretrain a 1.5-billion-parameter transformer (roughly GPT-2 XL scale, a reasonable size for a first serious Indic-language base model). They write standard PyTorch DistributedDataParallel training code, mixed precision, Adam optimizer. The job crashes with an out-of-memory error before a single activation tensor is allocated. Nothing about the model architecture changed between "works on paper" and "OOM in practice." What changed is that they never accounted for what Adam actually costs in bytes, and they never questioned whether every GPU truly needs a full redundant copy of everything. This chapter works out exactly why that job dies at 24 GB, and how DeepSpeed's ZeRO optimizer (Zero Redundancy Optimizer) reclaims that memory by refusing to store the same bytes four times over.

Where the bytes actually go

During mixed-precision training with Adam, four categories of state live on the GPU for every parameter: the fp16 parameters used for the forward and backward matrix multiplies (2 bytes per parameter), the fp16 gradients produced by backward (2 bytes), and the optimizer's own bookkeeping, kept in fp32 for numerical stability: a master fp32 copy of every parameter (4 bytes), the first-moment estimate m (4 bytes), and the second-moment estimate v (4 bytes). That fp32 optimizer bundle is 12 bytes per parameter for Adam specifically; Rajbhandari, Rasley, Ruwase and He call this constant K in their 2020 paper "ZeRO: Memory Optimizations Toward Training Trillion Parameter Models" (SC20), where K = 12 for Adam. Add it up: 2 + 2 + 12 = 16 bytes of guaranteed model-state memory per parameter, before a single activation is stored.

For our team's model, Ψ = 1.5 × 10⁹ parameters, so baseline memory per GPU is 16Ψ = 24 × 10⁹ bytes = 24 GB. On a 24 GB A10G, that is the entire card, with nothing left for activations, the optimizer's internal scratch space, or any batch at all. Crucially, this number does not shrink as more GPUs are added: standard data parallelism replicates the full model state onto every rank, so a cluster of 4 GPUs and a cluster of 400 GPUs both carry 24 GB of dead weight per card. That redundancy, not the compute, is what ZeRO attacks.

Three stages of partitioning the redundancy away

ZeRO keeps data parallelism's execution model (every GPU processes a different data shard, computes its own forward and backward) but stops replicating storage. Across N data-parallel ranks it partitions the model state into three optional, cumulative stages:

Stage 1, P_os: only the fp32 optimizer states (the 12Ψ bytes) are split into N equal shards, one per GPU. Parameters and gradients stay fully replicated. Memory per GPU: 4Ψ + 12Ψ/N.

Stage 2, P_os+g: gradients are additionally partitioned. Each GPU only ever materializes the gradient shard it needs to update its own slice of the optimizer state; the rest are reduced into other ranks and discarded locally. Memory per GPU: 2Ψ + (2Ψ + 12Ψ)/N.

Stage 3, P_os+g+p: parameters are partitioned too. Every GPU permanently owns only 1/N of the fp16 parameters; the full tensor for a given layer is reconstructed on demand via a collective communication call right before that layer is used, then freed immediately after. Memory per GPU: (2Ψ + 2Ψ + 12Ψ)/N = 16Ψ/N, a fully linear reduction in N.

Plugging in the team's numbers, Ψ = 1.5 × 10⁹, N = 4:

def zero_stage_memory(psi, Nd, stage, K=12):
    if stage == 0:      # baseline DDP: everything replicated
        total = 2*psi + 2*psi + K*psi
    elif stage == 1:    # partition optimizer states only
        total = 2*psi + 2*psi + K*psi / Nd
    elif stage == 2:    # partition optimizer states + gradients
        total = 2*psi + (2*psi + K*psi) / Nd
    elif stage == 3:    # partition everything
        total = (2*psi + 2*psi + K*psi) / Nd
    else:
        raise ValueError("stage must be 0-3")
    return total / 1e9  # bytes -> GB

for stage in [0, 1, 2, 3]:
    print(stage, round(zero_stage_memory(1.5e9, 4, stage), 2))

Tracing it by hand confirms the printed output before running anything: stage 0 gives 2(1.5e9) + 2(1.5e9) + 12(1.5e9) = 3e9 + 3e9 + 18e9 = 24e9 → 24.0. Stage 1 gives 3e9 + 3e9 + 18e9/4 = 3e9 + 3e9 + 4.5e9 = 10.5e9 → 10.5. Stage 2 gives 3e9 + (3e9 + 18e9)/4 = 3e9 + 5.25e9 = 8.25e9 → 8.25. Stage 3 gives 24e9/4 = 6e9 → 6.0. The loop therefore prints exactly: 0 24.0, 1 10.5, 2 8.25, 3 6.0. Stage 3 alone drops the team from a hard OOM at 24 GB to 6 GB, leaving 18 GB of headroom on each 24 GB card for activations and a real batch size. The diagram below shows these four footprints to scale, with the replicated portion of each bar drawn solid and the partitioned portion drawn lighter, to make visible exactly which piece of state stopped being duplicated at each stage.

What ZeRO does not save you: communication

Memory is not free to reduce. Standard data parallelism uses ring all-reduce on the gradient tensor: each of N ranks sends and receives roughly 2Ψ(N−1)/N elements, which for any reasonably sized N is well approximated as 2Ψ total traffic per rank (about Ψ to reduce-scatter the sum, about Ψ to all-gather the result back to every rank). ZeRO-1 and ZeRO-2 reproduce this exactly: gradients are reduce-scattered once (~Ψ of traffic) so each rank ends up with only the shard it needs to update its slice of the optimizer state, and after the local Adam step, the updated fp16 parameters are all-gathered back to every rank (~Ψ more). Total: ~2Ψ, identical to plain DDP. This is the paper's central claim for stages 1 and 2: memory drops by up to 8x as N grows large (2Ψ + 14Ψ/N approaches the 2Ψ floor as N → ∞, against the 16Ψ baseline), and the paper uses N = 64 as a worked example approaching that ceiling, for zero extra communication.

Stage 3 breaks that equivalence. Because parameters themselves are sharded, a given transformer layer's full weight tensor does not exist anywhere until it is all-gathered right before that layer runs, then freed. That all-gather happens once during the forward pass and again during the backward pass (the layer's parameters are needed a second time to compute its local gradient), adding roughly Ψ of traffic on top of the gradient reduce-scatter, so total communication becomes approximately 3Ψ instead of 2Ψ: a 1.5x tax versus stages 1 and 2. On a communication-bound cluster (slower interconnect, many nodes), that 50% increase can translate directly into wall-clock slowdown even though every byte of GPU memory stage 3 frees is real.

The misconception to retire

The mistake students most often make here is assuming ZeRO works like tensor or pipeline parallelism, where each GPU is permanently responsible for computing only a slice of the model. It is not. ZeRO is still data parallelism: every GPU runs the forward and backward pass for the entire model, on its own batch shard. What ZeRO partitions is storage, not computation. A stage-3 GPU that owns 1/4 of a given layer's weights does not compute 1/4 of that layer's matrix multiply; it briefly reconstructs the full weight tensor via an all-gather, runs the full computation like any other rank, and then discards the borrowed 3/4 the moment it is done. This is why ZeRO's memory savings come with a communication cost and not a compute-parallelism benefit: nothing about it reduces FLOPs per GPU, unlike sharding the model across ranks the way tensor parallelism does.

Offloading the optimizer entirely: ZeRO-Offload

Partitioning across GPUs is one axis of relief; DeepSpeed has an orthogonal one, moving state off the GPU altogether. ZeRO-Offload (Jie Ren, Samyam Rajbhandari, et al., "ZeRO-Offload: Democratizing Billion-Scale Model Training," USENIX ATC 2021) keeps fp16 parameters and the forward/backward compute on the GPU, but moves the entire fp32 optimizer bundle, the master weights, m, and v, into CPU RAM, and runs the Adam update itself on the CPU. Every step, gradients (2Ψ bytes) must cross PCIe from GPU to CPU, and the freshly updated fp16 parameters (2Ψ bytes) must cross back, a total of 4Ψ bytes of PCIe traffic per step. For this worked example, picture the same team having graduated to a single rented A100 for the PCIe/throughput arithmetic below (the A100's 40/80 GB capacity and 312 TFLOPS dense-fp16 peak are what the following numbers assume; the opening 24 GB OOM scenario was on A10G cards, a different, cheaper part). For the 1.5B-parameter model, 4Ψ = 6 × 10⁹ bytes = 6 GB per step. Assume PCIe Gen4 x16 sustaining roughly 25 GB/s of effective throughput per direction (below the ~31.5 GB/s theoretical peak, accounting for protocol overhead): the 3 GB grad-down and 3 GB param-up transfers can happen in parallel over the link's two directions, so transfer time is approximately 3 GB / 25 GB/s ≈ 0.12 s = 120 ms per step. Compare that against compute time using the standard approximation that one forward-plus-backward pass costs about 6 FLOPs per parameter per token (Kaplan et al., 2020, "Scaling Laws for Neural Language Models"): for a batch of 2048 tokens, FLOPs ≈ 6 × 1.5 × 10⁹ × 2048 ≈ 1.84 × 10¹³. An A100 sustaining roughly 150 TFLOPS in practice on dense fp16 matmuls (well below its 312 TFLOPS peak, a realistic utilization figure for well-optimized transformer training) completes that in 1.84 × 10¹³ / 1.5 × 10¹⁴ ≈ 0.123 s = 123 ms. The PCIe transfer (~120 ms) and the compute (~123 ms) land almost exactly on top of each other. That is not a coincidence to celebrate; it means the offload transfer is only free if DeepSpeed can overlap it with computation on the next step, which is precisely the engineering trick ZeRO-Offload implements: it pipelines the CPU-side Adam update against the GPU's next forward pass rather than treating the transfer as a blocking step. Without that overlap, offload would roughly double wall-clock step time for a model at this size and batch.

ZeRO-Infinity (Rajbhandari, Ruwase, Rasley, Smith, He, "ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning," SC21, 2021) extends the same idea one tier further, to NVMe SSD, for models whose optimizer state does not even fit in CPU RAM. NVMe bandwidth per drive is typically a few GB/s, well below the tens of GB/s available between CPU RAM and GPU, so ZeRO-Infinity leans even harder on prefetching and bandwidth-aware partitioning to hide that slower tier behind compute; the further down the memory hierarchy state gets pushed, the more overlap engineering is required to keep the GPU fed.

All of this is configured, not hand-coded, in DeepSpeed. A stage-2 job with CPU-offloaded optimizer state looks like this:

{
  "train_batch_size": 32,
  "fp16": { "enabled": true },
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": { "device": "cpu", "pin_memory": true },
    "allgather_bucket_size": 5e8,
    "reduce_bucket_size": 5e8
  },
  "optimizer": {
    "type": "AdamW",
    "params": { "lr": 3e-4, "betas": [0.9, 0.999], "eps": 1e-8 }
  }
}

zero_optimization.stage: 2 selects P_os+g from the formulas above. offload_optimizer.device: "cpu" layers ZeRO-Offload's CPU Adam on top of that partitioning; pin_memory: true asks the OS for page-locked host memory so PCIe transfers can use direct memory access instead of an extra CPU-side copy, which matters for hitting the bandwidth figures used above. allgather_bucket_size and reduce_bucket_size chunk the collective communication into buckets of that many elements so DeepSpeed can start transferring one bucket while the backward pass is still computing gradients for earlier layers, the software mechanism that makes the overlap argument above actually hold in practice rather than remaining a napkin calculation.

Active recall

Attempt each question before reading its answer.

1. In one sentence each, what does ZeRO-1, ZeRO-2, and ZeRO-3 partition that the previous stage did not?
2. A team trains a Ψ = 3 × 10⁹ parameter model with Adam across N = 8 GPUs, ZeRO stage 2. Compute the per-GPU memory in GB.
3. The original 1.5B-parameter, 4-GPU example is redone with N = 8 GPUs instead of 4. Recompute all four numbers (baseline, stage 1, stage 2, stage 3) and state which one changes least in absolute terms and why.
4. Explain, using the reduce-scatter/all-gather picture, why ZeRO-3's communication volume is about 1.5x that of ZeRO-1/ZeRO-2, not equal to it.
5. Redo the ZeRO-Offload PCIe timing example assuming an older PCIe Gen3 x8 link with an effective 12 GB/s per direction instead of 25 GB/s. Does the transfer still hide behind the ~123 ms compute time?
6. True or false, with justification: "Since ZeRO-3 gives the largest memory reduction, it should always be the default choice over ZeRO-2."

Answer 1. ZeRO-1 (P_os) partitions only the fp32 optimizer states (master weights, momentum, variance) across the N ranks. ZeRO-2 (P_os+g) additionally partitions the fp16 gradients, so each rank only ever holds the gradient shard needed for its slice of the optimizer update. ZeRO-3 (P_os+g+p) additionally partitions the fp16 parameters themselves, reconstructing the full tensor for a layer on demand via all-gather and freeing it right after use.

Answer 2. Stage 2 formula: 2Ψ + (2Ψ + 12Ψ)/N. With Ψ = 3 × 10⁹, N = 8: 2(3e9) + (2(3e9) + 12(3e9))/8 = 6e9 + (6e9 + 36e9)/8 = 6e9 + 42e9/8 = 6e9 + 5.25e9 = 11.25e9 bytes = 11.25 GB.

Answer 3. Baseline is unchanged at 24.0 GB, since standard DDP replicates the full state on every rank regardless of N; N never appears in that formula. Stage 1: 4Ψ + 12Ψ/8 = 6e9 + 18e9/8 = 6e9 + 2.25e9 = 8.25 GB (down from 10.5). Stage 2: 2Ψ + 14Ψ/8 = 3e9 + 21e9/8 = 3e9 + 2.625e9 = 5.625 GB (down from 8.25). Stage 3: 16Ψ/8 = 24e9/8 = 3.0 GB (down from 6.0, exactly halved, since it is the only stage with a purely linear 1/N dependence and no fixed replicated term). The baseline changes least, in fact not at all, because it has no N-dependence to begin with; this is exactly the redundancy ZeRO exists to eliminate.

Answer 4. ZeRO-1 and ZeRO-2 only ever move gradients: one reduce-scatter to sum and distribute gradient shards (~Ψ traffic) and one all-gather to redistribute the updated parameters after the local optimizer step (~Ψ traffic), for ~2Ψ total, matching plain data-parallel all-reduce. ZeRO-3 additionally shards the parameters at rest, so the full weight tensor for each layer must be reassembled via an extra all-gather immediately before that layer's forward computation, and reassembled again before its backward computation (to compute that layer's local gradient). That second reconstruction adds roughly one more Ψ of all-gather traffic that stages 1 and 2 never needed, bringing the total to ~3Ψ, a 1.5x increase over the ~2Ψ baseline.

Answer 5. The two transfers are still 3 GB grad-down and 3 GB param-up. At 12 GB/s effective per direction, each takes 3 GB / 12 GB/s = 0.25 s = 250 ms, and since the two directions can run concurrently on a full-duplex link, the offload step still takes about 250 ms (versus 123 ms of compute). Communication no longer hides behind computation, it now exceeds it by roughly 2x, so without deeper pipelining across multiple micro-batches, this configuration would slow the step down rather than being effectively free, unlike the Gen4 case where the two were roughly matched.

Answer 6. False. Stage 3 gives the deepest memory reduction, but it does so at roughly 1.5x the communication volume of stages 1 and 2, because of the extra parameter all-gathers described above. If a model already fits comfortably under stage 2, moving to stage 3 only adds communication overhead for memory savings you did not need, which can slow training on interconnect-limited clusters. The right choice is the smallest stage that gets the job within the memory budget, exactly the sequence worked through in this chapter: try stage 1, then stage 2, and reach for stage 3 (and offload beyond it) only when the model genuinely does not fit otherwise.

Think About It

Think about this: How would you explain deepspeed zero: extreme memory efficiency 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 deepspeed zero: extreme memory efficiency, 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.