企业官网建设流程全解析

• 推理引擎
• 大模型

本文以 DeepSpeed 官方博客《An Order-of-Magnitude Larger and Faster Training with ZeRO-2》为核心骨架，结合本仓库内 DeepSpeed 子模块的完整源码（deepspeed/runtime/zero/、deepspeed/runtime/activation_checkpointing/），系统讲解 ZeRO-2 如何在原始 ZeRO（优化器状态分区）基础上，进一步消解梯度、激活与碎片化内存开销，从而支撑 170B 参数级别模型的高效训练。读者读完将掌握 ZeRO-2 的四大内存优化维度、zero_optimization与activation_checkpointing配置项的完整语义，以及DeepSpeedZeroOptimizer的底层实现机制。

一、ZeRO-2 是什么：从优化器状态到全谱系内存优化

原始 ZeRO（Zero Redundancy Optimizer，零冗余优化器）的核心洞察是：分布式数据并行训练中，优化器状态（如 Adam 的 momentum 与 variance）在每张 GPU 上被冗余保存，造成了大量不必要的显存占用。ZeRO 的第一步（Stage 1）通过将优化器状态按数据并行度切分、每张卡只持有并更新自己那一份分区，消除了这份冗余。

ZeRO-2 则将这一思想从"优化器状态"扩展到了训练过程中的完整内存消耗谱系。官方文档明确指出，ZeRO-2 在原始 ZeRO 的优化器状态内存优化之外，新增了三项核心技术：

1. 梯度的内存消减——梯度不再每张卡冗余保存，而是像优化器状态一样按数据并行度切分；
2. 激活内存的优化——通过激活检查点（Activation Checkpointing）等手段压缩前向传播保存的中间激活；
3. 内存碎片的消除——通过连续内存分配器预分配大块缓冲区并支持碎片整理，避免训练中大量小张量的创建与释放把显存"打碎"。

文档宣称的综合效果是：将深度学习训练的规模与速度提升一个数量级（order of magnitude），允许训练高达 1700 亿（170B）参数的模型，速度相比当时的最先进方案（state of the art）提升最多 10 倍。需要说明的是，这里的"数量级""10 倍"是官方文档给出的项目自述结论，实际加速比会随模型、硬件与配置不同而变化。

二、ZeRO-2 的四大内存优化维度

2.1 优化器状态分区（继承自 ZeRO Stage 1）

训练大模型时，以 Adam 优化器 + 混合精度为例，每个模型参数在显存中实际占据 16 份空间：fp16 权重 2 份 + fp32 权重 4 份 + momentum 4 份 + variance 4 份 + 梯度 2 份。其中 optimizer 状态（fp32 权重 + momentum + variance）占比最大，且在数据并行下每张卡都存一份完整副本。

ZeRO Stage 1 将这 16 份优化器相关状态按数据并行度 N 切分，每张卡只持有 1/N。在DeepSpeedZeroOptimizer的初始化代码中可以看到对应实现：参数组被展平（flatten）后按数据并行度切成近似相等的分区，每个进程只从自己的分区克隆出 fp32 master 权重：

# deepspeed/runtime/zero/stage_1_and_2.py（节选） # divide the flat weights into near equal partition equal to the data parallel degree # each process will compute on a different part of the partition data_parallel_partitions = self.get_data_parallel_partitions( self.bit16_groups_flat[i], i) self.parallel_partitioned_bit16_groups.append(data_parallel_partitions) # A partition of the fp32 master weights that will be updated by this process. self.single_partition_of_fp32_groups.append( self.parallel_partitioned_bit16_groups[i][partition_id].to( self.device).clone().float().detach())

对应配置stage: 1。这一步完成后，本地优化器只包含属于自己分区的那部分参数，因此本地优化器的状态（momentum、variance 等）也只与自己的分区相关（源码注释原话："the local optimizer only saves the states(momentum, variance, etc.) related to its partition's params"）。

2.2 梯度分区（ZeRO-2 的核心新增）

ZeRO-2 的 Stage 2 在 Stage 1 基础上，把梯度也按数据并行度切分。反向传播时，每张卡并不保存全部梯度，而只保留自己分区对应的那一份梯度，之后用它更新自己的优化器状态分区。

在源码中，Stage 1 与 Stage 2 的区分仅由一个布尔量控制：

# ZeRO stage 1 (False) or 2 (True) self.partition_gradients = partition_grads

以及配置枚举中的明确注释：

# deepspeed/runtime/zero/config.py class ZeroStageEnum(int, Enum): """ Enum class for possible zero stages """ disabled = 0 optimizer_states = 1 # 仅优化器状态分区 gradients = 2 # 优化器状态 + 梯度分区 weights = 3 # 优化器状态 + 梯度 + 权重分区

梯度分区通过create_reduce_and_remove_grad_hooks()在参数上注册反向钩子实现：每个梯度计算完成后立即按桶（bucket）进行 reduce-scatter 归约，归约后该进程只持有自己分区那部分的平均梯度，其余部分被释放。这避免了全量梯度在显存中的滞留。

2.3 激活内存优化

ZeRO-2 的另一大内存来源是前向传播保存的激活值——对长序列、大 batch 的 Transformer 模型，激活值内存常常超过模型权重本身。ZeRO-2 通过激活检查点（Activation Checkpointing）技术，只保留少量"检查点"激活用于反向重计算，从而大幅降低激活内存峰值。

在本仓库中，激活检查点的配置独立于 ZeRO 配置，位于deepspeed/runtime/activation_checkpointing/config.py：

"activation_checkpointing": { "partitioned_activations": [true|false], "number_checkpoints": 100, "contiguous_memory_optimization": [true|false], "cpu_checkpointing": [true|false], "profile": [true|false], "synchronize_checkpoint_boundary": [true|false] }

各配置项语义（均来自DeepSpeedActivationCheckpointingConfig源码）：

• partition_activations（默认False）：是否将激活值按数据并行度分区存储，进一步降低单卡激活内存；
• number_checkpoints（默认None）：设置激活检查点的数量，控制内存与重计算开销的权衡；
• contiguous_memory_optimization（默认False）：是否将激活存储到连续内存缓冲区，减少碎片并提升内存利用率；
• cpu_checkpointing（默认False）：是否将激活检查点卸载到 CPU 内存；
• profile（默认False）：是否输出激活检查点相关内存画像；
• synchronize_checkpoint_boundary（默认False）：是否在检查点边界做同步。

2.4 内存碎片消除

深度学习训练中，前向/反向会不断创建、销毁各种大小的临时张量（激活、梯度、中间结果），长此以往显存被"打碎"成大量不连续的小空闲块，即使总空闲空间足够也无法分配大张量。ZeRO-2 引入了**连续内存分配器（Contiguous Memory Allocator）**应对这一问题。

其实现位于deepspeed/runtime/zero/contiguous_memory_allocator.py的ContiguousMemoryAllocator类：初始化时一次性预分配一块大缓冲区（self.buffer = torch.zeros(size, ...)），后续张量分配都从这块缓冲区中"切"出；当剩余空间不足以容纳请求的大小时，分配器会自动执行碎片整理（defragmentation），把散落的小空闲块合并成连续大块，并重新绑定参数张量的数据指针（_reset_param_data）。这一机制与配置中的contiguous_gradients相互配合：梯度被复制进连续缓冲区，避免反向传播期间的碎片化。

三、ZeRO-2 配置实战：zero_optimization完整参数详解

DeepSpeedZeroConfig（deepspeed/runtime/zero/config.py）对 ZeRO-2 相关的全部配置项做了精确定义。以下是在 JSON 配置中启用 ZeRO-2（Stage 2）的完整形态：

"zero_optimization": { "stage": 2, "contiguous_gradients": true, "reduce_scatter": true, "reduce_bucket_size": 500000000, "allgather_partitions": true, "allgather_bucket_size": 500000000, "overlap_comm": false, "load_from_fp32_weights": true, "offload_optimizer": { "device": "cpu", "pin_memory": true }, "round_robin_gradients": false, "ignore_unused_parameters": true }

各参数语义与源码依据如下：

⚠️ 注意两点限制：

1. ZeRO-2 是数据并行内的优化，需要至少 2 张 GPU 才有分区收益；单卡场景应选择 ZeRO-Offload（offload_optimizer）或 ZeRO-3。
2. 启用reduce_scatter时，通信数据类型仅支持 fp16/bf16（源码断言：ZeRO-2 supports only float16 or bfloat16 communication_data_type with reduce scatter enabled），且gradient_predivide_factor必须为 1.0。

四、源码级剖析：DeepSpeedZeroOptimizer的工作机制

ZeRO-1/2 的统一实现类是DeepSpeedZeroOptimizer（deepspeed/runtime/zero/stage_1_and_2.py，约 2491 行）。其核心流程可以概括为四个阶段：

4.1 初始化：展平、分区与 fp32 master 权重

初始化时，每个参数组先被移到 CPU 以释放显存空间创建 flat buffer，随后：

1. 展平：调用flatten_dense_tensors_aligned将组内参数展平为连续张量，并按 NCCL 4 字节对齐要求（nccl_start_alignment_factor = 2，即 fp16 元素数的 2 倍）对齐；
2. 分区：get_data_parallel_partitions将扁平权重切成与数据并行度相等的 N 份，每张卡持有 1/N；
3. fp32 master 权重：每张卡从自己那份 fp16 分区克隆出 fp32 master 权重（clone().float().detach()），替换本地优化器的参数，此后本地优化器只对自己分区做状态维护。

同时源码用move_to_cpu与see_memory_usage反复记录内存占用，为每步内存分配做诊断（partition_count、partition_size等信息会在各 rank 打印）。

4.2 反向传播：桶式梯度归约

梯度分区通过反向钩子实现：create_reduce_and_remove_grad_hooks在每个参数上注册钩子，梯度一旦就绪就按reduce_bucket_size聚合到桶（IPG bucket）中，桶满即发起 reduce-scatter，归约完成后各卡只保留自己分区那份平均梯度，其余立即释放（_release_ipg_buffers）。overlap_comm开启时，归约在独立的reduction_stream上执行，与反向计算重叠。

4.3 优化器步骤：每卡更新自己的分区

由于本地优化器已被替换为只含本卡分区的 fp32 参数（见param_group['params'] = [self.single_partition_of_fp32_groups[i]]），优化器 step 时每张卡独立更新自己那份状态，全程无需其他卡的优化器状态——这正是"零冗余"的本质。

4.4 每步结束：allgather 还原完整参数

更新完成后，通过allgather_partitions（allgather 集合）或一串 broadcast 把各卡的分区权重汇聚回完整参数，恢复模型权重供下一轮前向使用。allgather_bucket_size控制单次 allgather 的批大小，避免为超大模型一次性申请过大通信缓冲。

4.5 CPU 卸载扩展：ZeRO-2 Offload

offload_optimizer将 Stage 1/2 的优化器状态与计算卸载到 CPU，此时contiguous_gradients被强制为 True（源码：self.contiguous_gradients = contiguous_gradients or cpu_offload），配合DeepSpeedCPUAdam在 CPU 上执行优化器更新，并使用 pin-memory 缓冲（temp_grad_buffer_for_cpu_offload = torch.zeros(...).pin_memory()）加速梯度传输。round_robin_gradients进一步将梯度按轮转方式分配给各 rank 拷贝，避免单一 rank 成为 CPU 卸载的瓶颈。

五、ZeRO-2 与 ZeRO-3 的边界：何时用 Stage 2

从ZeroStageEnum可以清楚看到 ZeRO 三个阶段的递进关系：

• Stage 1：只分区优化器状态；
• Stage 2（ZeRO-2）：分区优化器状态 + 梯度；
• Stage 3（ZeRO-3）：再进一步分区模型权重，参数按需 allgather（配合stage3_max_live_parameters、stage3_max_reuse_distance、stage3_prefetch_bucket_size等参数控制驻留参数数量）。

ZeRO-2 不需要像 Stage 3 那样在前向/反向中反复 gather 权重，通信开销更低，适合"显存紧张但未到单卡装不下权重"的规模；当模型权重本身都无法放入单卡时，则需要切换到 Stage 3（甚至配合offload_param卸载权重到 CPU/NVMe，即 ZeRO-Infinity 路线）。此外，overlap_comm的动态默认值也印证了这一设计：源码 validator 中，只有 Stage 3 才会默认开启通信重叠。

六、验证与上手路径

• 本仓库的 DeepSpeed 子模块位于 benchmark/third_party/DeepSpeed，ZeRO-2 核心实现见 stage_1_and_2.py，配置解析见 config.py；
• 激活检查点配置与解析见 activation_checkpointing/config.py，连续内存分配器实现见 contiguous_memory_allocator.py；
• DeepSpeed 官方在文档中建议参考 Megatron 教程中 GPT 系列模型的训练示例来上手 ZeRO-2（对应本仓库 docs/_tutorials/megatron.md 教程）；更详细的技术原理可查阅官方技术报告（原文档所引 arxiv 论文）。

配置示例可直接作为 DeepSpeed 训练脚本的ds_config.json使用，再通过deepspeed.initialize(config=ds_config)或命令行--deepspeed_config传入即可启用 Stage 2 训练。对于 170B 级别的模型训练，通常还需结合张量并行（Megatron 式模型并行）与流水线并行，将 ZeRO-2 作为数据并行维度上的内存优化手段。

七、小结

ZeRO-2 的本质，是把"零冗余"原则从优化器状态推广到梯度，并叠加激活检查点与连续内存分配两项工程优化，全面压制训练过程中的显存开销。它让每个数据并行 rank 只维护自己分区的优化器状态与梯度，通过 reduce-scatter 归约 + allgather 还原的通信模式，以少量额外通信换取显存占用的大幅下降——这正是官方文档所称"训练规模与速度提升一个数量级"的技术根基。理解这四大优化维度及其在DeepSpeedZeroOptimizer中的落地方式，是正确配置 ZeRO-2、评估其收益与通信开销的前提。

• 推理引擎
• 大模型