Documentation Index
Fetch the complete documentation index at: https://docs.apiyi.com/llms.txt
Use this file to discover all available pages before exploring further.
概念定义
ZeRO（Zero Redundancy Optimizer）是一种革命性的内存优化技术，通过在数据并行进程间分片存储模型状态（优化器状态、梯度和参数），消除冗余存储，从而大幅降低每个GPU的内存占用。
详细解释
传统数据并行训练中，每个GPU都需要存储完整的模型参数、梯度和优化器状态，造成严重的内存冗余。以Adam优化器为例，每个参数需要存储：参数本身（FP16，2字节）、梯度（FP16，2字节）、动量（FP32，4字节）和二阶动量（FP32，4字节），总计12字节——是模型参数大小的6倍。
ZeRO通过将这些状态在多个GPU间分片存储，每个GPU只保存部分状态，需要时通过通信获取其他部分。这种方法在保持数据并行计算效率的同时，将内存占用降低了数个数量级，使得在有限的GPU资源上训练超大规模模型成为可能。
工作原理
ZeRO的三个阶段
Stage 1：优化器状态分片（Pos）
# 内存节省：4倍（对于Adam）
# 每个GPU只存储1/N的优化器状态
# 通信开销：无额外开销
optimizer_memory = model_size * 8 / gpu_count
Stage 2：梯度分片（Pos+g）
# 内存节省：8倍
# 每个GPU只保留需要更新的梯度部分
# 通信开销：与数据并行相同
gradient_memory = model_size * 2 / gpu_count
Stage 3：参数分片（Pos+g+p）
# 内存节省：与GPU数量成正比
# 参数也被分片，需要时通过All-Gather获取
# 通信开销：增加50%
parameter_memory = model_size * 2 / gpu_count
内存计算公式
每GPU内存需求 = (模型参数 × 系数) / GPU数量
其中系数：
- 标准数据并行：16倍（混合精度）
- ZeRO-1：12倍
- ZeRO-2：4倍
- ZeRO-3：2倍
ZeRO-Offload扩展
# CPU卸载配置
zero_config = {
"stage": 3,
"offload_optimizer": {
"device": "cpu",
"pin_memory": True
},
"offload_param": {
"device": "cpu",
"pin_memory": True
}
}
实际应用
2025年典型配置
- 单GPU微调40B+模型：ZeRO-3 + CPU Offload
- 8×A100训练175B模型：ZeRO-3 + NVMe Offload
- 512GPU训练2T参数模型：ZeRO-3 + 3D并行
DeepSpeed集成示例
from deepspeed import DeepSpeedConfig
# ZeRO-3配置
ds_config = {
"zero_optimization": {
"stage": 3,
"overlap_comm": True,
"contiguous_gradients": True,
"sub_group_size": 1e9,
"reduce_bucket_size": 1e6,
"stage3_prefetch_bucket_size": 1e6,
"stage3_param_persistence_threshold": 1e5,
"stage3_max_live_parameters": 1e9,
"stage3_max_reuse_distance": 1e9,
"stage3_gather_16bit_weights_on_model_save": True
}
}
# 初始化模型
model_engine, optimizer, _, _ = deepspeed.initialize(
model=model,
config=ds_config
)
性能优化技巧
- 通信重叠：计算与通信并行，隐藏通信延迟
- 梯度累积：减少通信频率
- 参数预取：提前获取下一步需要的参数
- 混合精度：配合BF16/FP8进一步降低内存
扩展性分析
- 线性扩展：内存占用随GPU数量线性降低
- 通信开销：ZeRO-3增加约1.5倍通信量
- 带宽要求：建议使用InfiniBand或高速NVLink
- CPU内存：Offload需要充足的系统内存
相关概念
延伸阅读
