DeepSpeed ZeRO优化策略：显存管理与大规模训练实践
一、大模型训练的显存瓶颈：参数量的指数级增长
大语言模型的参数量已从亿级增长到千亿级，训练过程中的显存消耗成为首要瓶颈。以175B参数的模型为例，仅模型参数（FP32）就需要700GB显存，加上梯度、优化器状态（Adam需要两倍参数量的动量和方差），总显存需求超过2.8TB。即便使用A100 80GB GPU，也需要数十张卡才能容纳一个模型实例。
传统的数据并行（Data Parallelism）将完整模型复制到每张GPU上，显存效率极低；模型并行（Model Parallelism）将模型切分到多张GPU，但通信开销大且扩展性差。DeepSpeed ZeRO（Zero Redundancy Optimizer）通过消除数据并行中的冗余显存占用，在不增加通信开销的前提下大幅降低单卡显存需求。
本文将深入剖析ZeRO的三级优化策略，覆盖显存切分原理、通信机制和工程实践。
二、ZeRO优化策略原理
2.1 显存占用分析
训练过程中，GPU显存主要被四部分占用：模型参数（Parameters）、梯度（Gradients）、优化器状态（Optimizer States）和激活值（Activations）。以参数量为Ψ的模型为例，使用Adam优化器时的显存占用如下：
graph TB
subgraph "显存占用分解"
A["模型参数<br/>2Ψ (FP16) + 4Ψ (FP32主权重)"]
B["梯度<br/>2Ψ (FP16)"]
C["优化器状态<br/>4Ψ (动量) + 4Ψ (方差)"]
D["激活值<br/>与batch_size和序列长度相关"]
end
subgraph "ZeRO优化级别"
E["ZeRO-1: 切分优化器状态<br/>节省4x显存"]
F["ZeRO-2: +切分梯度<br/>节省8x显存"]
G["ZeRO-3: +切分参数<br/>节省N_d倍显存"]
end
A –> G
B –> F
C –> E
2.2 ZeRO三级优化详解
ZeRO-1：优化器状态切分
将优化器状态均匀切分到N_d张GPU上，每张GPU只存储1/N_d的优化器状态。在参数更新时，每张GPU只更新自己负责的参数分区，然后通过All-Gather同步更新后的完整参数。
class ZeRO1Optimizer:
"""ZeRO-1优化器简化实现"""
def __init__(self, model_params, optimizer, world_size, rank):
self.world_size = world_size
self.rank = rank
self.optimizer = optimizer
# 将参数按world_size切分
self.param_partitions = self._partition_params(model_params)
# 每个rank只维护自己分区的优化器状态
self.local_partition = self.param_partitions[self.rank]
def step(self):
# 1. 每个rank只更新自己分区的参数
for param in self.local_partition:
if param.grad is not None:
self.optimizer.step(param)
# 2. All-Gather同步更新后的完整参数
dist.all_gather(
tensor_list=self.param_partitions,
tensor=self.local_partition
)
def _partition_params(self, params):
"""将参数列表均匀切分为world_size个分区"""
param_list = list(params)
total = len(param_list)
partition_size = math.ceil(total / self.world_size)
partitions = []
for i in range(self.world_size):
start = i * partition_size
end = min(start + partition_size, total)
partitions.append(param_list[start:end])
return partitions
ZeRO-2：梯度切分
在ZeRO-1基础上，进一步将梯度切分。每张GPU只存储自己参数分区对应的梯度，通过Reduce-Scatter操作实现梯度的聚合和切分同步完成。
class ZeRO2Optimizer(ZeRO1Optimizer):
"""ZeRO-2优化器简化实现"""
def backward(self, loss):
# 1. 反向传播计算梯度
loss.backward()
# 2. Reduce-Scatter：聚合梯度并切分
for param in self.model_params:
if param.grad is not None:
# 将完整梯度Reduce-Scatter到各rank
dist.reduce_scatter(
output=self.local_grad_partition,
input_list=param.grad
)
# 3. 非本地分区的梯度可以释放
for param in self.model_params:
if param not in self.local_partition:
param.grad = None # 释放非本地梯度
ZeRO-3：参数切分
ZeRO-3将模型参数也进行切分，每张GPU只存储1/N_d的参数。在前向和反向传播时，通过All-Gather临时获取需要的参数，计算完成后立即释放。
class ZeRO3ForwardHook:
"""ZeRO-3前向传播Hook：按需获取参数"""
def __init__(self, param_partition, rank, world_size):
self.param_partition = param_partition
self.rank = rank
self.world_size = world_size
def pre_forward(self, module, input):
"""前向传播前：All-Gather获取完整参数"""
for name, param in module.named_parameters():
# 从各rank收集完整参数
full_param = torch.empty_like(param.data)
dist.all_gather(
tensor_list=full_param,
tensor=self._get_local_partition(param)
)
# 临时替换为完整参数
param.data = full_param
def post_forward(self, module, input, output):
"""前向传播后：释放非本地参数"""
for name, param in module.named_parameters():
if not self._is_local(param):
# 释放非本地分区参数
param.data = self._get_local_partition(param)
三、DeepSpeed工程实践
3.1 配置文件设计
DeepSpeed通过JSON配置文件控制ZeRO的优化级别和行为。
{
"train_batch_size": 256,
"gradient_accumulation_steps": 4,
"fp16": {
"enabled": true,
"loss_scale": 0,
"initial_scale_power": 16
},
"zero_optimization": {
"stage": 3,
"offload_param": {
"device": "cpu",
"pin_memory": true
},
"offload_optimizer": {
"device": "cpu",
"pin_memory": true
},
"overlap_comm": true,
"contiguous_gradients": true,
"sub_group_size": 1e9,
"reduce_bucket_size": 5e8,
"stage3_prefetch_bucket_size": 5e8,
"stage3_param_persistence_threshold": 1e5
}
}
3.2 通信优化策略
ZeRO-3引入了额外的通信开销（前向和反向传播中的All-Gather），需要通过通信与计算重叠来隐藏延迟。
class CommunicationOverlap:
"""通信与计算重叠优化"""
def __init__(self, model, world_size):
self.model = model
self.world_size = world_size
self.prefetch_queue = asyncio.Queue(maxsize=2)
async def prefetch_params(self, layer_idx: int):
"""异步预取下一层参数"""
next_layer = self.model.get_layer(layer_idx + 1)
if next_layer is not None:
# 在当前层计算时，异步获取下一层参数
await self.prefetch_queue.put(
self._all_gather_params(next_layer)
)
def forward_with_overlap(self, x):
"""带通信重叠的前向传播"""
for idx, layer in enumerate(self.model.layers):
# 等待当前层参数就绪
params = self.prefetch_queue.get_nowait() \\
if not self.prefetch_queue.empty() \\
else self._all_gather_params(layer)
# 执行当前层计算
x = layer(x, params)
# 异步预取下一层参数
asyncio.create_task(self.prefetch_params(idx))
return x
四、架构权衡与边界分析
4.1 通信开销与显存节省的权衡
ZeRO-1的通信量与纯数据并行相同（仅增加一次All-Gather），但显存节省有限；ZeRO-3的显存节省最大，但通信量增加约1.5倍。对于通信带宽受限的集群，ZeRO-2可能是更优选择——在显存节省和通信开销之间取得平衡。
4.2 CPU Offload的延迟代价
将优化器状态和参数卸载到CPU可以进一步降低GPU显存需求，但CPU-GPU之间的数据搬运会显著增加训练时间。实测中，CPU Offload可能导致训练速度下降30%-50%。建议仅在GPU显存确实不足时启用Offload，且优先卸载优化器状态（访问频率低），而非参数（访问频率高）。
4.3 激活值重计算的取舍
ZeRO主要优化参数、梯度和优化器状态的显存占用，但激活值（Activations）的显存消耗随batch_size和序列长度线性增长。结合激活值重计算（Activation Checkpointing），只保留关键层的激活值，其余层在反向传播时重新计算，可以进一步降低显存需求，代价是增加约30%的计算量。
五、总结
DeepSpeed ZeRO通过三级优化策略，从优化器状态、梯度到参数逐步消除数据并行中的冗余显存占用。ZeRO-1通信开销最小，适合通信受限场景；ZeRO-2在显存和通信之间取得平衡；ZeRO-3显存节省最大，但需要通信与计算重叠来隐藏延迟。
落地建议：从ZeRO-1开始验证训练流程的正确性，逐步升级到ZeRO-2和ZeRO-3；启用CPU Offload前先评估训练速度的下降是否可接受；结合激活值重计算进一步降低显存需求，但需权衡额外的计算开销。
