LLM Systems 18: Memory Optimization in Distributed Training

1. Memory Consumption for LLM training

DDP 与 MP 回顾：

• DP
  • 优点：
    • Good compute/communication efficiency
  • 缺点：
    • Poor memory efficiency(Every device has one copy of model)
• MP
  • 优点：
    • Good memory efficiency
  • 缺点：
    • Poor compute/communication efficiency(5% of peak perf in training 40B model with Megatron)

Memory Consumption in DDP (Mixed Precision)

设模型参数量为 \(N\)，每个 GPU 需要内存为 \(20N + \text{act}\)

• 模型参数：\(2N\) bytes
• 激活值：\(b \times s \times h \times l\)
• 反向传播梯度：\(2N\) bytes
• 优化器状态（以 Adam 为例）：\(N \times 4 \times 4\) ，包括：
  • 参数、梯度、动量、方差
• 示意图如下（DP=2）：

Other Memory Usages in DDP

• Temporary Buffers:
  • Storing intermediate results. Operations such as gradient all-reduce, or gradient norm computation tend to fuse all the gradients into a single flattened buffer before applying the operation in an effort to improve throughput.
• Memory Fragmentation

Common Approaches to Reduce Memory

• ZeRO: Partition Optimizer States, Gradients, Parameters
• Reducing Activation Memory
  • Activation Checkpoint, Compression
• CPU Offload
• Memory Efficient Optimizer
  • Maintaining coarser-grained statistics of model parameters and gradients

2. Partitioning and reducing the memory for data parallel training

ZeRO - Zero Redundancy Optimizer

• Key idea:
  • Eliminating data redundancy in DDP by partitioning the optimizer states (zero-1), gradients (zero-2), parameters (zero-3)
• Widely used for large language model training.
  • 7B model memory: 120GB ➔ 30GB (with 4GPUs)
• Implemented in DeepSpeed.

2.1 Partition optimizer states

示意图如下：

• 切分优化器状态
• 通信过程：
  • backward 时 ReduceScatter 获取 Reduce 后梯度分片；
  • 更新本地分片参数；
  • AllGather 获取全部更新后参数；

2.2 Partition gradients

• Key idea:
  • Each GPU compute all parameter gradients for its data partition
  • but only stores one partition of gradients instead of all gradients.
  • Passing the gradients out of its responsibility to the GPU responsible for those gradients.
• Memory for gradients reduced by K times. (K=# of GPUs)

示意图：

• In backward pass, GPU 0,1,2 hold temporary buffers for the gradients that GPU 3 is responsible for (M3)
• GPU 0,1,2 pass the M3 gradients to GPU 3 (ncclReduce)
  • bucket（partition）级别的 ncclReduce，等价于全部 gradients 级别的 ReduceScatter，通信量不变。
• Then GPU0, GPU1, GPU2 delete M3 gradients, GPU 3 keeps M3 gradients.
• Continue backward pass: GPU 0,1,3 hold temporary buffers for the gradients that GPU 2 is responsible for (M2)
• ...

2.3 Partition parameters

示意图：

• Partition parameters to K parts
• During forward, GPU0 send first params to other GPUs(ncclBroadcast)
• compute forward for first part layers
• GPU1,2,3 delete first part parameters, GPU0 still keep first part params
• Continue forward for next part, GPU1 send its params to other GPUs(ncclBroadcast), compute forward for second part layers
• ...
• During backward, compute backward for 4th part layers. Reduce grads to GPU3 (same as ZeRO stage 2)
• GPU0,1,2 delete 4th part parameters and gradients
• Continue backward, GPU2 broadcast its parameters to GPU0, 1, 3. All GPUs perform backward to calculate local gradients for part 3
• ...

Memory Consumption & Communication Cost for ZeRO

阶段	显存消耗公式	通信开销变化
Original DDP	\(4N + MN\)	基线（\(2N\)，ReduceScatter+AllGather 各 \(N\)）
ZeRO-1	\(4N + \frac{MN}{K}\)	无额外通信总量（与 DDP 相同）
ZeRO-2	\(2N + \frac{(2+M)N}{K}\)	无额外通信总量（与 DDP 相同）
ZeRO-3	\(\frac{(4+M)N}{K}\)	\(3N\)（fwd/bwd 需要广播模型参数，加上梯度 ReduceScatter，共 \(3N\)）

3. Other memory optimizations

• Partitioned Activation Checkpointing
  • 结合 MP 切分激活值
• Constant Size Buffers (similar to Bucketing in PyTorch DDP)
  • Buffer is used in doing all-reduce to improve bandwidth.
  • Modern implementations fuses all the parameters into a single buffer.
  • ZeRO uses constant size buffers to be more efficient for a large model.
• Memory Defragmentation
  • Long-lived memory (Model parameters, Optimizer state): Store together
  • Short-lived memory (Discarded activations)
• We can further improve by memory reuse as in LightSeq.

Reduce Communication in ZeRO DDP

• Applying block-wise quantization technique to parameters during forward (broadcast params). FP16 ➔ INT 8.
  • zeropoint quantization
• Apply quantization during backward ReduceScatter, FP16➔ INT8 or INT4
• better partition parameters: maintain a full set of parameters on each node

Combining ZeRO with Pipeline Parallelism and Tensor Parallelism (3D Parallelism)

• PP 切分模型不同层
• 单 Node 内进行 TP 切分
• 利用 ZeRO 作用于 DP 维度
  • 在不同 DP Rank 中处于相同位置（即处理相同模型层、相同张量分片）的 GPU 之间（上图用带 "ZeRO" 标识的双向箭头连接），建立 ZeRO 优化组，对其对应的优化器状态、梯度及参数进行分摊切片，从而大幅削减 DP 维度的显存冗余。

4. Conclusion

• ZeRO : reducing memory footprint in distributed training
  • ➔ enables training significantly larger models
• Key idea: partition optimizer states, gradients, and parameters.
• Pros:
  • Lower memory usages significantly.
  • Scalable, flexible, easy-to-use.
• Cons:
  • Some stages introduce extra communication overhead, depending on infrastructure (PCI-E / NVLink)