ZeRO State Sharding Principles

When training Large Language Models using standard Data Parallelism, the system replicates the entire model, gradients, and optimizer states across every available GPU. For an 8 billion parameter model like meta-llama/Meta-Llama-3.1-8B, full parameter training requires approximately 16 to 18 bytes of memory per parameter when using mixed-precision training and the AdamW optimizer. This translates to roughly 128 GB of VRAM just to hold the baseline training states. If you have a server with four 80GB A100 GPUs, standard Data Parallelism copies this 128 GB requirement to all four devices, resulting in immediate out-of-memory errors across the board. To fit training jobs onto available hardware, engineers rely on the Zero Redundancy Optimizer, commonly referred to as ZeRO. ZeRO eliminates memory redundancies by partitioning training states across the distributed cluster rather than replicating them, allowing you to train massive models by pooling the memory of multiple GPUs.

The Memory Baseline and State Sharding Mechanics

To understand how ZeRO fragments memory, you first need to break down the 16 bytes of overhead required for every single parameter during training.

1. Model Weights: 2 bytes per parameter (FP16 or BF16 precision).
2. Gradients: 2 bytes per parameter (FP16 or BF16 precision).
3. Optimizer States: 12 bytes per parameter for the AdamW optimizer. This includes a 32-bit floating point copy of the weights (4 bytes), a first momentum vector (4 bytes), and a second momentum variance vector (4 bytes).

The progression of memory sharding topologies. Each successive ZeRO stage reduces static memory per GPU at the cost of increased network communication bandwidth.

Rather than storing 100% of these components on every device, ZeRO divides them across the total number of available GPUs in the cluster. This division happens in three distinct stages, allowing infrastructure operators to balance memory savings against network communication overhead.

ZeRO Stage 1 Partitioning Optimizer States

The largest single consumer of memory during fine-tuning is the optimizer state, taking up 12 of the 16 bytes. ZeRO Stage 1 targets this bottleneck by keeping the model weights and gradients fully replicated on every GPU, but splitting the optimizer states into $N$ equal chunks, where $N$ is the number of GPUs.

$Memory_{ZeRO1} = Parameters \times \left(4 + \frac{12}{N}\right)$

If you are fine-tuning meta-llama/Meta-Llama-3.1-8B on a cluster of four 80GB A100 GPUs, the calculation is straightforward. The weights and gradients require 4 bytes per parameter, while the 12-byte optimizer state is divided by 4.

$8 \times 10^9 \times \left(4 + \frac{12}{4}\right) = 8 \text{ billion} \times 7 \text{ bytes} = 56 \text{ GB}$

By applying ZeRO Stage 1, the memory footprint per GPU drops from 128 GB to 56 GB. This easily fits inside an 80GB A100 accelerator, leaving 24 GB of headroom for forward activations, batching, and context length.

ZeRO Stage 2 Adding Gradient Sharding

ZeRO Stage 2 extends the partitioning strategy by sharding the gradients alongside the optimizer states. Each GPU only holds the gradients for the specific parameters it is assigned to update.

$Memory_{ZeRO2} = Parameters \times \left(2 + \frac{14}{N}\right)$

Using the same meta-llama/Meta-Llama-3.1-8B model on the four-GPU A100 cluster, the static memory footprint drops further.

$8 \times 10^9 \times \left(2 + \frac{14}{4}\right) = 8 \text{ billion} \times 5.5 \text{ bytes} = 44 \text{ GB}$

This aggressive reduction means that training an 8B model is now feasible on a dense workstation packed with four 48GB RTX 6000 Ada generation cards, assuming a small batch size and heavily optimized activation memory.

ZeRO Stage 3 and Fully Sharded Data Parallelism

For much larger architectures like Qwen/Qwen2.5-32B, even ZeRO Stage 2 is insufficient. A 32-billion parameter model requires over 500 GB of baseline memory. ZeRO Stage 3 resolves this by sharding everything. The weights, gradients, and optimizer states are all distributed equally across the hardware.

$Memory_{ZeRO3} = \frac{Parameters \times 16}{N}$

During execution, GPUs broadcast the specific model layers they need on-demand over the interconnect, performing matrix operations, and immediately discarding the fetched weights to free up VRAM. PyTorch provides a native implementation of this exact mechanic called Fully Sharded Data Parallelism (FSDP).

If you deploy Qwen/Qwen2.5-32B across an eight-GPU node, the footprint per GPU becomes highly manageable.

$Memory_{ZeRO3} = \frac{32 \text{B} \times 16}{8} = 64 \text{ GB per GPU}$

Memory footprint scaling for Qwen2.5-32B across an 8-GPU cluster. ZeRO Stage 3 ensures the 64 GB requirement fits safely inside 80GB accelerators.

Interconnect Hardware and Communication Overhead

The primary trade-off of ZeRO Stage 3 is the massive tax it places on the system interconnect. Because GPUs must continuously pull missing weights from neighboring cards during the forward and backward passes, the network bandwidth becomes a hard physical bottleneck.

If you attempt to run ZeRO Stage 3 on an 8-GPU RTX 4090 rig connected via standard PCIe Gen 4 lanes (yielding roughly 32 GB/s per link), the compute cores will sit idle while waiting for weights to transfer over the motherboard. This causes severe drops in training throughput. In contrast, enterprise accelerators like the H100 utilize NVLink Gen 4, providing 900 GB/s of bidirectional bandwidth. This high-speed interconnect effectively masks the ZeRO Stage 3 network overhead, keeping the Tensor Cores saturated.

Implementing ZeRO Offloading Configurations

When physical GPU VRAM is completely exhausted, you can configure ZeRO to offload optimizer states to the host CPU RAM via the PCIe bus. This allows massive models to train on smaller hardware footprints, though the CPU data transfers will significantly increase the time it takes to complete an epoch.

You can configure these limits natively using a DeepSpeed JSON configuration file passed to standard Hugging Face training scripts.

{
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "offload_param": {
      "device": "none"
    },
    "overlap_comm": true,
    "contiguous_gradients": true
  },
  "fp16": {
    "enabled": true
  }
}

In this setup, the stage: 3 directive shards all parameters, while offload_optimizer pushes the heavy 12-byte AdamW states entirely into system memory. If you run this workload on an Apple Silicon machine like a Mac Studio with 128GB of unified memory, the CPU offload penalty is virtually eliminated because the GPU and CPU inherently share the exact same physical memory blocks, making it an excellent development environment for validating distributed configurations locally.