Data Parallelism

Premium

Understanding communication primitives and DDP's gradient synchronization mechanism

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

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\Phi$ bytes
Gradients: $2\Phi$ bytes

fp32 part (optimizer):

Parameter copy: $4\Phi$ bytes (used for accurate updates)
First moment $m$ : $4\Phi$ bytes (exponential moving average of the gradient)
Second moment $v$ : $4\Phi$ bytes (moving average of the squared gradient)

Total: $2\Phi + 2\Phi + 4\Phi + 4\Phi + 4\Phi = 16\Phi$ bytes

Component	Precision	Memory (bytes)
Parameters	fp16	$2\Phi$
Gradients	fp16	$2\Phi$
Parameter copy	fp32	$4\Phi$
First moment $m$	fp32	$4\Phi$
Second moment $v$	fp32	$4\Phi$
Total		$16\Phi$

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 $m$ and $v$ needed?

Adam does not simply "take a step along the gradient." It needs to maintain two cross-step historical statistics: $m$ (the exponential moving average of the gradient, equivalent to momentum) and $v$ (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 $m$ and $v$ are "the accumulated experience from all past steps."