LogoCookLLM Docs
LogoCookLLM Docs
HomeCookLLM

Principles

Tokenization
Tokenization BasicsBPE AlgorithmGPT TokenizersBPE Training Engineering
Model Architecture
Transformer LM
From token ids to logitsEmbedding and LM Head
Attention Mechanisms
From Self-Attention to GQAAttention Sink
Position Encoding
Position Encoding BasicsRoPE Math DerivationRoPE ImplementationLength Extrapolation
GPU Programming Basics
GPU Architecture BasicsTensor LayoutTriton Basics: Vector Add
FlashAttention
Flash Attention PrinciplesFrom Naive Implementation to Auto-TuningBlock Pointers and Multi-Dim SupportCausal Masking OptimizationGrouped Query AttentionBackward Pass Implementation
Distributed Training
Data ParallelismZeRO OptimizerFully Sharded Data ParallelTensor ParallelismPipeline ParallelismMulti-Dimensional Hybrid Parallelism

Hands-on Training

Overview
Pretraining
Pretraining DataTokenizer TrainingModel ArchitectureData PipelineTraining LoopMonitoring and Validation
X (Twitter)
SystemsDistributed Training

ZeRO Optimizer

Premium

Progressive de-redundancy: three-stage sharding from optimizer states to parameters

Get code access

In the previous chapter we saw DDP's memory problem: to guarantee training consistency (by synchronizing gradients via All-Reduce), each GPU needs to store the complete model state. Four GPUs means 4 complete copies (parameters, gradients, optimizer states). The core idea of ZeRO (Zero Redundancy Optimizer) is straightforward: since the final state is consistent anyway, let each GPU store only a portion, and communicate to fetch the rest when needed.

Redundancy Analysis of Training State

Let's first quantify DDP's waste. Taking mixed precision + Adam as an example, with NNN GPUs training a model with Φ\PhiΦ parameters, each GPU needs to store:

Log in to continue reading

This is premium content. Please log in to access the full article.

  • Parameters (fp16): 2Φ2\Phi2Φ bytes
  • Gradients (fp16): 2Φ2\Phi2Φ bytes
  • Optimizer states (fp32): 12Φ12\Phi12Φ bytes (parameter copy + first moment + second moment)

For a total of 16Φ16\Phi16Φ bytes, of which optimizer states account for 75%.

NNN GPUs means NNN-fold redundancy: global storage of 16NΦ16N\Phi16NΦ bytes, when only 16Φ16\Phi16Φ bytes are actually needed. ZeRO's three stages eliminate these redundancies one by one, in order from largest to smallest.

Data Parallelism

Understanding communication primitives and DDP's gradient synchronization mechanism

Fully Sharded Data Parallel

Understanding FSDP's Intra-Tensor sharding and All-Gather/Reduce-Scatter communication patterns

Table of Contents

Redundancy Analysis of Training State
ZeRO Stage 1: Sharding Optimizer States
Parameter Assignment Strategy
Gradient Synchronization
Training Loop Comparison
ZeRO Stage 2: Sharding Gradients
ZeRO Stage 3: Sharding Parameters
Parameter Sharding
Communication Pattern
Communication Overhead Comparison
ZeRO-3's Sharding Method: Inter-Tensor
Summary