Training large language models introduces a memory challenge that linear scaling of hardware cannot solve alone. When using standard Distributed Data Parallel (DDP), the training process hits a hard ceiling defined by the VRAM of a single GPU. Understanding the exact composition of this memory footprint is necessary for engineering systems capable of training models with billions or trillions of parameters.

The Anatomy of Model Memory

To optimize memory usage, we must first quantify the cost of training a single parameter. It is a common misconception that a model with parameter count $\Psi$ requires $4\Psi$ bytes (assuming 32-bit floats) or $2\Psi$ bytes (assuming 16-bit floats) of memory. In practice, the memory consumption during training with the Adam optimizer and mixed precision is significantly higher.

In a standard mixed-precision training pipeline (using FP16 or BF16 for computation and FP32 for weight updates), the system must maintain several copies of the model state. For every individual parameter in the model, the memory allocation consists of:

Model Parameters (FP16/BF16): The weights used for the forward and backward passes. Requires 2 bytes.
Gradients (FP16/BF16): The computed gradients from the backward pass. Requires 2 bytes.
Optimizer States (FP32): The Adam optimizer requires high precision to maintain stability. It stores:
- Master Weights: A high-precision copy of the parameters. Requires 4 bytes.
- Momentum ( $m_t$ ): First moment estimate. Requires 4 bytes.
- Variance ( $v_t$ ): Second moment estimate. Requires 4 bytes.

Summing these components yields the memory constant for mixed-precision training:

$M_{param} = 2 + 2 + 4 + 4 + 4 = 16 \text{ bytes}$

Therefore, a model with $\Psi$ parameters requires $16\Psi$ bytes of static memory. A 7 billion parameter model, often considered "small" in modern LLM contexts, requires $7 \times 10^9 \times 16$ bytes, or approximately 112 GB of VRAM, merely to load the weights and optimizer states. This exceeds the capacity of an NVIDIA A100 (80GB) before a single token is processed.

Redundancy in Distributed Data Parallel

DDP functions by replicating the entire model state across every worker in the cluster. If you deploy a cluster of $N$ GPUs, DDP creates $N$ identical copies of the model parameters, gradients, and optimizer states. The communication step in DDP (AllReduce) synchronizes gradients across workers, but it does not reduce the memory footprint on any individual device.

The efficiency of DDP degrades as model size increases. While DDP allows you to scale batch size by adding GPUs, it does not allow you to scale model size. The memory requirement per GPU remains constant regardless of the cluster size:

$\text{Memory}_{DDP} = 16\Psi + \text{Activation Memory} + \text{Fragmentation}$

This architecture results in massive memory redundancy. In a cluster with 16 GPUs training a 1 billion parameter model (16 GB footprint), the total cluster memory usage is $16 \times 16 \text{ GB} = 256 \text{ GB}$ . However, the unique information stored is only 16 GB. The remaining 240 GB is duplicate data.

Memory allocation breakdown per parameter in mixed-precision training.

ZeRO: Eliminating Redundancy

The Zero Redundancy Optimizer (ZeRO) addresses this inefficiency by acknowledging that while all GPUs need access to all weights during the forward and backward passes, they do not need to persist all weights, gradients, and optimizer states simultaneously.

ZeRO partitions (shards) the model states across the available data-parallel processes. If there are $N_d$ GPUs, ZeRO splits the data such that each GPU owns $1/N_d$ of the total state. This sharding can be applied in three progressive stages, each offering greater memory savings at the cost of increased communication complexity.

Stage 1: Optimizer State Sharding

The optimizer states (Master Weights, Momentum, Variance) constitute the bulk of the memory footprint (12 bytes out of 16). In Stage 1, these states are sharded across $N_d$ GPUs. Each GPU updates only its assigned partition of the optimizer states.

$\text{Memory}_{Stage1} = 2\Psi \text{ (Weights)} + 2\Psi \text{ (Grads)} + \frac{12\Psi}{N_d} \text{ (Opt States)}$

Stage 2: Gradient Sharding

Stage 2 extends sharding to the gradients. As gradients are computed during the backward pass, they are reduced and sharded immediately, rather than aggregated locally.

$\text{Memory}_{Stage2} = 2\Psi \text{ (Weights)} + \frac{2\Psi + 12\Psi}{N_d}$

Stage 3: Parameter Sharding

Stage 3 is the core of FSDP. It shards the model parameters themselves. At this stage, a GPU only persists a fraction of the model. When a specific layer is needed for computation, the parameters are gathered from other GPUs, used, and then immediately discarded to free memory.

$\text{Memory}_{Stage3} = \frac{16\Psi}{N_d}$

The theoretical limit of ZeRO Stage 3 allows the memory footprint per device to approach zero as the number of devices $N_d$ increases, leaving the majority of VRAM available for activations and larger batch sizes.

Comparative Analysis: DDP vs. FSDP Scaling

The divergence in memory efficiency becomes pronounced as we scale the number of GPUs. With DDP, adding GPUs yields no reduction in memory pressure per device. With FSDP (ZeRO Stage 3), memory pressure decreases linearly with the addition of hardware.

For instance, a scenario training a Large Language Model with $\Psi$ parameters. The chart below demonstrates the maximum model size trainable on an 80GB A100 GPU as the cluster size increases.

Maximum trainable model size (in Billions of parameters) per GPU as cluster size scales.

In the DDP configuration (red line), the maximum model size is strictly capped around 3.5 billion parameters per GPU (allowing buffer for activations). Adding 60 more GPUs does not change this limit. In the FSDP configuration (blue line), the capacity scales linearly. With 64 GPUs, the cluster can effectively train a model approaching 175 billion parameters, as the $16\Psi$ static state is distributed thinly across the cluster.

Impact on Activation Memory

It is important to note that ZeRO only reduces the memory footprint of the model state. It does not inherently reduce the memory required for activations, the intermediate outputs of layers stored for the backward pass. Activation memory depends on batch size, sequence length, and transformer architecture (e.g., hidden dimension, attention heads).

While FSDP frees up massive amounts of VRAM from model parameters, training terabyte-scale models often requires combining FSDP with Activation Checkpointing (recomputing activations during the backward pass) to keep activation memory within bounds. We will implement this integration in Chapter 3.

By shifting from DDP to FSDP, we move from a regime where model architecture is constrained by single-device limits to a regime where model size is constrained only by total cluster capacity and network bandwidth.

Was this section helpful?