尽管 PyTorch 的 DistributedDataParallel (DDP) 中实现的标准数据并行对于提升吞吐量很有效,但它引入了显著的内存瓶颈。数据并行组中的每个 GPU 都维护着模型参数、梯度和优化器状态的完整副本。对于大型模型,特别是那些使用 Adam 等优化器(会存储动量和方差)的模型,仅优化器状态所需的内存就可能超过模型参数大小的两倍或更多。这种复制在 GPU 计算能力达到饱和之前很久就成了限制因素。Microsoft 的 DeepSpeed 库通过其零冗余优化器 (ZeRO) 直接解决了这种内存效率低下问题。ZeRO 不会复制整个训练状态,而是将其划分到可用数据并行设备上。这使得您可以训练比单个 GPU 所能容纳的模型大许多倍的模型,而无需依赖于纯粹的模型并行或流水线并行的复杂性。内存冗余问题为了充分理解 ZeRO 的作用,让我们细分标准数据并行训练期间单个 GPU 上的内存消耗。总内存 $M$ 可以建模为三个主要部分的和:$$ M = M_P + M_G + M_O $$各部分为:$M_P$ 是模型参数的内存。$M_G$ 是梯度的内存,其大小通常与 $M_P$ 相同。$M_O$ 是优化器状态的内存。对于 Adam 这样的优化器,这包括动量和方差,通常导致 $M_O \ge 2 \times M_P$。在包含 $K$ 个 GPU 的标准 DDP 配置中,每个 GPU 都持有这三个部分的完整副本,导致大量的冗余。digraph G { rankdir=TB; splines=false; node [shape=record, style="filled", fillcolor="#dee2e6", fontname="Arial"]; subgraph cluster_0 { label = "标准数据并行 (DDP)"; bgcolor="#f8f9fa"; node [fillcolor="#a5d8ff"]; gpu0 [label="{GPU 0 | {参数 (P) | 梯度 (G) | 优化器 (O)}}"]; gpu1 [label="{GPU 1 | {参数 (P) | 梯度 (G) | 优化器 (O)}}"]; gpuK [label="{GPU K | {参数 (P) | 梯度 (G) | 优化器 (O)}}"]; } gpu0 -> gpu1 [style=invis]; gpu1 -> gpuK [style=invis]; }标准数据并行中的内存布局,每个 GPU 都持有模型参数、梯度和优化器状态的完整副本。ZeRO:训练状态划分ZeRO 通过在 GPU 之间分片模型和优化器状态来系统地消除这种冗余。它分三个渐进阶段实现,让您选择最适合需求的优化级别。ZeRO 阶段 1 (ZeRO-1)第一阶段侧重于 Adam 等优化器最主要的内存冗余来源:优化器状态。ZeRO-1 在数据并行进程之间划分优化器状态。每个 GPU 现在仅负责更新其分配的参数分区。在优化器步骤中,all-gather 操作会收集所有更新的参数分片,以确保每个 GPU 都有一个完整且最新的模型,用于下一个前向传播。划分: 优化器状态 ($M_O$)。内存节省: 显著,尤其适用于 Adam/AdamW。ZeRO 阶段 2 (ZeRO-2)ZeRO-2 在阶段 1 的基础上,进一步划分梯度。在反向传播期间,不再是每个 GPU 持有完整梯度集然后通过 all-reduce 操作对其进行平均,而是使用 reduce-scatter 操作。此操作同时计算平均值并将分片结果分散到对应的 GPU。这避免了同时存储完整梯度和参数所导致的瞬时内存峰值。划分: 优化器状态 ($M_O$) + 梯度 ($M_G$)。内存节省: 大量。消除梯度冗余。ZeRO 阶段 3 (ZeRO-3)这是最先进的阶段,使得训练真正庞大的模型成为可能。ZeRO-3 划分所有内容:优化器状态、梯度和模型参数本身。每个 GPU 仅持有一个整个模型的切片。在前向和反向传播期间,ZeRO-3 仅在计算需要时才在每个 GPU 上动态重构模型的完整层。一个层的参数在执行之前会从所有参与的 GPU 收集,并在之后立即回收内存。这意味着在任何给定时间,内存峰值使用量与单个层的大小成比例,而不是与整个模型成比例。划分: 优化器状态 ($M_O$) + 梯度 ($M_G$) + 参数 ($M_P$)。内存节省: 最大。使得支持万亿参数的模型成为可能。digraph G { rankdir=TB; splines=false; node [shape=record, style="filled", fillcolor="#dee2e6", fontname="Arial"]; subgraph cluster_1 { label = "ZeRO-3 数据并行"; bgcolor="#f8f9fa"; subgraph cluster_g0 { label="GPU 0"; bgcolor="#e9ecef"; p0 [label="P_0", shape=box, style="filled", fillcolor="#b2f2bb"]; g0 [label="G_0", shape=box, style="filled", fillcolor="#a5d8ff"]; o0 [label="O_0", shape=box, style="filled", fillcolor="#ffec99"]; } subgraph cluster_g1 { label="GPU 1"; bgcolor="#e9ecef"; p1 [label="P_1", shape=box, style="filled", fillcolor="#b2f2bb"]; g1 [label="G_1", shape=box, style="filled", fillcolor="#a5d8ff"]; o1 [label="O_1", shape=box, style="filled", fillcolor="#ffec99"]; } subgraph cluster_gK { label="GPU K"; bgcolor="#e9ecef"; pk [label="P_K", shape=box, style="filled", fillcolor="#b2f2bb"]; gk [label="G_K", shape=box, style="filled", fillcolor="#a5d8ff"]; ok [label="O_K", shape=box, style="filled", fillcolor="#ffec99"]; } } p0 -> g0 -> o0 [style=invis]; p1 -> g1 -> o1 [style=invis]; pk -> gk -> ok [style=invis]; }ZeRO-3 的内存布局:参数 (P)、梯度 (G) 和优化器状态 (O) 全部划分到可用 GPU 上。ZeRO-Offload:将内存扩展到 CPU 和 NVMe即使使用 ZeRO-3,一个庞大模型的总状态大小也可能超过集群中所有可用 GPU 内存的总和。ZeRO-Offload 通过将训练状态的某些部分移动到更多、但速度较慢的内存层级来扩展划分层级。CPU 卸载: 您可以将 DeepSpeed 配置为将划分的优化器状态,甚至参数,卸载到主机 CPU 的内存中。在训练步骤中,所需数据会从 CPU 移动到 GPU 进行计算,然后移回。这对于优化器步骤尤其有效,该步骤通常计算量小但内存需求大。NVMe 卸载: 对于最极端的情况,DeepSpeed 可以将参数和优化器分区卸载到快速 NVMe 固态硬盘。这提供了可用内存的极大扩展,使得支持万亿参数模型成为可能,但代价是与 CPU 内存相比,I/O 延迟显著更高。这种卸载能力让大型模型训练变得普及,使得在 GPU 显存有限但拥有大量系统内存或快速存储的系统上也能实现训练。DeepSpeed 的实际实现将 DeepSpeed 集成到 PyTorch 训练脚本中非常直接。主要改变包括初始化 DeepSpeed 引擎以及修改训练循环以使用 DeepSpeed 的方法。首先,您创建一个 ds_config.json 文件。此文件是所有 DeepSpeed 功能的控制面板。ZeRO-2 带 CPU 卸载的 ds_config.json 示例:{ "train_batch_size": 16, "train_micro_batch_size_per_gpu": 2, "steps_per_print": 100, "optimizer": { "type": "AdamW", "params": { "lr": 0.0001, "betas": [0.9, 0.999], "eps": 1e-8, "weight_decay": 3e-7 } }, "zero_optimization": { "stage": 2, "offload_optimizer": { "device": "cpu", "pin_memory": true }, "contiguous_gradients": true, "overlap_comm": true }, "fp16": { "enabled": true } }接下来,您修改您的训练脚本。重点是 deepspeed.initialize 函数,它会封装您的模型和优化器。import torch import deepspeed # 假设模型和优化器已定义 # model = MyTransformerModel() # optimizer = torch.optim.AdamW(model.parameters()) # DeepSpeed 初始化 model_engine, optimizer, _, _ = deepspeed.initialize( model=model, optimizer=optimizer, config='ds_config.json' ) # 训练循环 for step, batch in enumerate(data_loader): # 将批次移动到正确的设备 batch = to_device(batch, model_engine.local_rank) # 前向传播 loss = model_engine(batch) # 反向传播 - 使用模型引擎 model_engine.backward(loss) # 优化器步骤 - 使用模型引擎 model_engine.step()注意训练循环中的变化:deepspeed.initialize 返回一个 model_engine,它会替换您原始的模型以进行主要操作。loss.backward() 调用被 model_engine.backward(loss) 替换。optimizer.step() 调用被 model_engine.step() 替换。DeepSpeed 根据您提供的配置处理所有底层分片、通信和卸载的复杂性。这种简洁的 API 使您只需更改 JSON 配置,即可尝试不同的 ZeRO 阶段和卸载策略,而无需更改您的核心训练逻辑。