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

Data Parallelism

Premium

Understanding communication primitives and DDP's gradient synchronization mechanism

Get code access

The first step in training a large model is to get multiple GPUs working together. Data Parallelism is the most intuitive approach: each GPU holds a complete copy of the model, processes different data independently, and finally aggregates the gradients. In this chapter, we first understand the memory bottleneck of a single GPU, then learn the basics of multi-GPU communication, and finally dive into the implementation of DDP.

Memory Composition of Single-GPU Training

Log in to continue reading

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

A 7B-parameter model needs only 14 GB for its fp16 weights (7B × 2 bytes), yet it cannot be trained on an 80GB A100. Where did all the memory go?

The answer is: during training, the GPU stores not only the parameters but also the gradients and optimizer states.

Memory Requirements of Mixed-Precision Training

Modern training commonly uses Mixed Precision: forward and backward passes are computed in fp16 (fast), but parameter updates use fp32 (accurate). Taking the Adam optimizer as an example, assume the model has Φ\PhiΦ parameters:

fp16 part (forward/backward):

  • Parameters: 2Φ2\Phi2Φ bytes
  • Gradients: 2Φ2\Phi2Φ bytes

fp32 part (optimizer):

  • Parameter copy: 4Φ4\Phi4Φ bytes (used for accurate updates)
  • First moment mmm: 4Φ4\Phi4Φ bytes (exponential moving average of the gradient)
  • Second moment vvv: 4Φ4\Phi4Φ bytes (moving average of the squared gradient)

Total: 2Φ+2Φ+4Φ+4Φ+4Φ=16Φ2\Phi + 2\Phi + 4\Phi + 4\Phi + 4\Phi = 16\Phi2Φ+2Φ+4Φ+4Φ+4Φ=16Φ bytes

ComponentPrecisionMemory (bytes)
Parametersfp162Φ2\Phi2Φ
Gradientsfp162Φ2\Phi2Φ
Parameter copyfp324Φ4\Phi4Φ
First moment mmmfp324Φ4\Phi4Φ
Second moment vv

Why is an fp32 parameter copy needed?

During training, both forward and backward passes are computed in fp16 (fast and memory-efficient), but fp16 has only about 3 significant digits of precision. When the learning rate is very small, the update term in param += learning_rate × gradient may be so tiny that fp16 rounds it directly to zero, and the model stops learning. So Adam internally maintains an fp32-precision copy of the parameters, performs the update in fp32, and then converts the result back to fp16 for the next forward pass.

Why are mmm and vvv needed?

Adam does not simply "take a step along the gradient." It needs to maintain two cross-step historical statistics: mmm (the exponential moving average of the gradient, equivalent to momentum) and vvv (the moving average of the squared gradient, used to adaptively adjust the learning rate for each parameter). The gradient is "which way to go this step," while and are "the accumulated experience from all past steps."

Actual memory for a 7B model:

  • Parameters: 14 GB
  • Gradients: 14 GB
  • Optimizer states: 84 GB (parameter copy 28GB + mmm 28GB + vvv 28GB)
  • Total: about 112 GB (not counting activations yet)

A single A100 has only 80 GB of memory, which cannot even hold the model's training state.

Training Memory per GPU

Mixed Precision + Adam
Total = 16Φ bytes
12Φ
2Φ
2Φ
Optimizer States(fp32 params + m + v)
Gradients(fp16)
Parameters(fp16)
Example: 7B Model (Φ = 7×10⁹)
Optimizer States84 GB
Gradients14 GB
Parameters14 GB
Total112 GB
A100: 80 GB

Distributed Training

From data parallelism to multi-dimensional hybrid parallelism — understanding the core parallel strategies of large model training

ZeRO Optimizer

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

Table of Contents

Memory Composition of Single-GPU Training
Memory Requirements of Mixed-Precision Training
Communication Primitives
Broadcast
All-Reduce
Reduce-Scatter
All-Gather
DataParallel: The Most Naive Multi-GPU Approach
How DDP Works
Ring All-Reduce: Efficient Gradient Synchronization
Gradient Synchronization Mechanism
Limitations of DDP
Summary
v
fp32
4Φ4\Phi4Φ
Total16Φ16\Phi16Φ
mmm
vvv