DeepSpeed ZeRO-3 零冗余优化实战指南：原理、配置、参数协调与调试

【免费下载链接】DeepSpeed DeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective. 【免费下载链接】DeepSpeed 项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed

本文以 DeepSpeed 的官方 API 文档 docs/code-docs/source/zero3.rst 为主线，结合仓库中的配置与实现源码（deepspeed/runtime/zero/config.py、deepspeed/runtime/zero/offload_config.py、deepspeed/utils/tensor_fragment.py 等），系统讲解 ZeRO 三阶段的分区原理、ZeRO-Infinity 卸载能力、完整的配置示例，以及如何构造超大模型、协调分区权重、调试与修改模型状态。读完本文，你将能独立写出一份可运行的 ZeRO-3 / ZeRO-Infinity 训练配置，并掌握分布式场景下访问与修改参数、梯度、优化器状态的正确姿势。

ZeRO 的核心思想：消除数据并行中的三类内存冗余

ZeRO（Zero Redundancy Optimizer）解决的是经典数据并行（Data Parallelism）中的根本性浪费：在传统数据并行下，每个进程都完整复制一份模型状态（model states）——即优化器状态（optimizer states）、梯度（gradients）和参数（parameters）。模型越大，这份复制带来的冗余越不可接受。

ZeRO 的做法是将这三类模型状态按数据并行维度进行分区（partition），每个进程只保存其中的一部分，而不是全部复制。通过这种方式，ZeRO 在保留数据并行计算粒度和通信效率的同时，大幅提升了内存利用效率。

ZeRO 分为三个递进的阶段：

  1. ZeRO Stage 1 —— 分区优化器状态：优化器状态（例如 Adam 优化器的 32 位权重、一阶动量 exp_avg、二阶动量 exp_avg_sq）被切分到各进程，每个进程只负责更新属于自己的那部分分区。
  2. ZeRO Stage 2 —— 追加分区梯度：用于更新模型权重的规约后的 16 位梯度同样被分区，每个进程只保留与其优化器状态分区相对应的梯度部分。
  3. ZeRO Stage 3 —— 追加分区参数：16 位模型参数也被分区到各进程。ZeRO-3 会在前向与反向传播过程中自动完成参数的收集（all-gather）与再分区（re-partition）。

在 deepspeed/runtime/zero/config.py 中，这一阶段体系被建模为 ZeroStageEnum：disabled = 0、optimizer_states = 1、gradients = 2、weights = 3、max_stage = 3，与文档描述一一对应。

从源码结构看，阶段划分本质上是一条递进的"卸载/分区链"：Stage 3 的分区实现（deepspeed/runtime/zero/parameter_offload.py 与 deepspeed/runtime/zero/stage3.py）复用了前两个阶段的梯度与优化器分区逻辑。可以将三个阶段理解为下表：

阶段分区优化器状态分区梯度分区参数可进一步 CPU/NVMe 卸载
Stage 1✔✘✘优化器状态可卸载（ZeRO-Offload 的雏形）
Stage 2✔✔✘优化器 + 梯度卸载（ZeRO-Offload）
Stage 3✔✔✔全部模型状态（ZeRO-Infinity）

从 ZeRO-Offload 到 ZeRO-Infinity

DeepSpeed 最早引入的卸载能力是 ZeRO-Offload：在 ZeRO-2 体系内，将优化器状态与梯度状态卸载到 CPU 内存，并在 CPU 上执行相应的优化器计算。而 ZeRO-Infinity 是面向 ZeRO-3 的下一代卸载能力：它继承了 ZeRO-Offload 的全部收益，并且能进一步卸载模型权重（参数），同时具备更有效的带宽利用以及计算与通信的相互重叠。

需要留意的是，原文档专门提醒：ZeRO-Infinity 与 ZeRO-Offload 与深度优化的 deepspeed.ops.adam.DeepSpeedCPUAdam 优化器配合效果最佳。推荐在 optimizer 配置段中声明优化器参数，交由 deepspeed.initialize() 来构建该优化器，而不是在模型代码中手工 new 一个，这样 CPU 端的优化器计算才能与卸载的数据布局完全对齐。

快速开始：把 ZeRO 跑起来

如果你尚未在 DeepSpeed 中训练过模型，可先参考入门教程 docs/_tutorials/getting-started.md。一旦你的模型已经能在 DeepSpeed 引擎下训练，启用 ZeRO-3 以及 ZeRO-Infinity 卸载就只需在 DeepSpeed 配置中增加 zero_optimization 一节——不需要改动模型代码（绝大多数模型）。

完整的 zero_optimization 选项清单可查阅仓库内的配置文档 docs/_pages/config-json.md；本文则聚焦 ZeRO-3 相关的用法与底层含义。

ZeRO 配置结构：三个 Pydantic 配置模型

所有 ZeRO 设置都由 zero_optimization 这一主配置键下的字典驱动，该字典会被解析并校验为 DeepSpeedZeroConfig 类（定义于 deepspeed/runtime/zero/config.py）。其中 offload_param 与 offload_optimizer 两个子配置分别由 DeepSpeedZeroOffloadParamConfig 与 DeepSpeedZeroOffloadOptimizerConfig 解析（两者定义于 deepspeed/runtime/zero/offload_config.py）。

由于这套配置类是自动文档化的数据模型，下面直接把各配置项的默认值与语义列出来（默认值均来自当前仓库源码，是撰写与调参的第一手依据）。

核心 ZeRO 配置项（DeepSpeedZeroConfig）

配置项（含别名）默认值含义
stage0（关闭）选择 ZeRO 阶段：0/1/2/3 分别对应关闭、分区优化器状态、追加分区梯度、追加分区参数
contiguous_gradientstrue梯度产生时拷贝到连续缓冲区，避免反向传播阶段的内存碎片化
reduce_scattertrue使用 reduce-scatter 而非 all-reduce 来平均梯度
reduce_bucket_size5e8 个元素一次规约/全规约的梯度元素数量，限制大模型下梯度聚合所需内存
allgather_partitionstrue使用 all-gather 集合通信而非一串 broadcast 来收集每步更新后的参数
allgather_bucket_size5e8 个元素一次 all-gather 的参数元素数量，限制大模型下的内存占用
overlap_comm动态（Stage 3 默认 true）尝试将梯度规约与反向计算重叠；见 overlap_comm_valid 校验器
stage3_prefetch_bucket_size5e7 个元素前向/反向中预取（fetch-ahead）的最大参数量
stage3_param_persistence_threshold1e5小于该阈值的参数不分区。阈值越小内存越省，但通信（尤其延迟敏感的小消息）会显著增加
stage3_model_persistence_thresholdsys.maxsize可常驻 GPU 不分区的参数总量上限，约束由上一项导致的不分区参数规模
stage3_max_live_parameters1e9每 GPU 在释放前最多常驻的参数数量；值越小越省内存但通信越多
stage3_max_reuse_distance1e9若某参数在该参数阈值内会被再次使用则不释放；值越小越省内存但通信越多
stage3_gather_16bit_weights_on_model_save（新名 gather_16bit_weights_on_model_save）false保存 16 位模型前先把分区权重聚合（consolidate），使 save_16bit_model() 落盘的是完整权重
sub_group_size1e9处理海量模型（万亿参数级）时的参数分块（tile）大小，供 ZeRO3-Offload 与 ZeRO-Infinity 使用
override_module_applytrue覆盖 nn.Module.apply 以在初始化时自动收集参数（详见下文 "Overriding Module.apply"）
ignore_unused_parameterstrue检测到未使用参数时是否报错终止；默认忽略，主要用于 Stage 2 动态网络
zero_hpz_partition_size1二级参数分区组中的 rank 数量
zero_quantized_weights / zero_quantized_gradients 等false对分区权重/梯度做量化以压缩通信
mics_shard_size / mics_hierarchical_params_gather-1 / falseMiCS 相关配置，见下文
leaf_module空需要被当作 ZeRO-3 叶模块处理的模块配置（DeepSpeedZeroLeafModuleConfig）
zenflownull是否启用 ZenFlow 优化

DeepSpeedZeroConfig 还带有若干模型校验器，从源码可以精确还原其约束逻辑：

  • overlap_comm_valid：当 overlap_comm 未显式给出时，若 stage == 3 则自动置为 true；
  • offload_ratio_check：offload_optimizer.ratio < 1.0 的部分卸载仅在 ZeRO Stage 3 下合法；
  • elastic_checkpoint_deprecated：Stage 3 的弹性检查点（elastic checkpointing）已废弃，应改用 Universal Checkpointing（见 deepspeed/checkpoint/universal_checkpoint.py）。

另外在配置模板字符串 ZERO_FORMAT（见 config.py）中还保留了历史兼容字段的完整清单，例如 cpu_offload、cpu_offload_param、cpu_offload_use_pin_memory。这些旧字段在 config.py 中被标记为 deprecated，迁移指南是：cpu_offload → offload_optimizer、cpu_offload_param → offload_param。新代码应直接使用新的子配置结构。

参数卸载配置（DeepSpeedZeroOffloadParamConfig）

配置项默认值含义
devicenone参数卸载目标：cpu 或 nvme；仅在 stage 3 下合法
nvme_pathnullNVMe 设备在文件系统中的挂载路径
buffer_count5参数 NVMe 卸载的缓冲池缓冲数量
buffer_size1e8每个缓冲的大小（字节数），构成缓冲池总容量
max_in_cpu1e9启用 NVMe 卸载时，仍保留在 CPU 内存中的参数元素数量上限
pin_memorytrue是否卸载到页锁定（pinned）CPU 内存；异步全带宽 GPU↔CPU 传输及卸载与计算重叠都依赖它。仅在宿主机 ulimit -l 紧张或内存极小时关闭

优化器卸载配置（DeepSpeedZeroOffloadOptimizerConfig）

配置项默认值含义
devicenone优化器状态卸载目标：cpu 或 nvme；对 stage 1/2/3 均合法。注意无论选什么设备，优化器计算都会被卸载到 CPU
nvme_pathnullNVMe 路径
buffer_count4优化器状态 NVMe 卸载缓冲池数量；至少应为优化器对每个参数维护的状态数（如 Adam 含参数、梯度、动量、方差 4 个状态）
pin_memorytrue是否使用页锁定内存
pipeline_readfalse基于 tile 的优化器 step 中，将下一 tile 的读取与当前 tile 的计算重叠（ZeRO-Infinity）
pipeline_writefalse将上一 tile 的写回与当前 tile 计算重叠
fast_initfalseNVMe 卸载时启用快速优化器初始化
ratio1.0卸载到 CPU Adam 的优化器状态比例（0~1）；< 1.0 即部分卸载，仅 stage 3 合法
super_offloadfalse面向 Superchip 的高性能 CPU 卸载（仅 stage 3）
cpuadam_cores_perc0.8super_offload=True 时 CPU Adam 使用的 CPU 核数比例

由浅入深的 ZeRO-3 配置示例

示例一：最简 ZeRO Stage 3

将优化器状态（stage 1）、梯度（stage 2）与参数（stage 3）全部进行分区。这是"纯 GPU 内存内分区"的基线配置：

{
    "zero_optimization": {
        "stage": 3
    },
    "fp16": {
        "enabled": true
    },
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": 0.001,
            "betas": [0.8, 0.999],
            "eps": 1e-8,
            "weight_decay": 3e-7
        }
    }
}

若使用 bf16 训练，可把 fp16 段替换为 "bf16": { "enabled": true }；是否启用取决于硬件与数值需求。需要说明的是，上面的 optimizer 段交给了 deepspeed.initialize() 去构建，这正是前文建议的用法。

示例二：把优化器状态和计算卸载到 CPU（ZeRO-Infinity 入门）

在上例基础上，通过 offload_optimizer.device = "cpu" 把 Adam 的 fp32 主权重与一、二阶动量卸载到 CPU，优化器更新步骤也在 CPU 上执行，从而腾出 GPU 显存：

{
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu"
        }
    }
}

示例三：进一步把参数也卸载到 CPU 内存

省下更多显存的做法是同时启用参数卸载。注意这与 ZeRO-Offload 的差别——ZeRO-3 能把模型权重本身也搬到 CPU：

{
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu"
        },
        "offload_param": {
            "device": "cpu"
        }
    }
}

这里有两种进阶玩法值得说明：

  • 部分优化器卸载：把 offload_optimizer.ratio 设为 0.0 ~ 1.0 之间的值，让一部分优化器状态留在 GPU、其余卸载到 CPU。源码校验器强制该特性仅用于 Stage 3（见上文 offload_ratio_check）。
  • 通信重叠：开启 offload_optimizer.pipeline_read / pipeline_write 可以重叠基于 tile 的优化器 step 的读写；这两项会影响训练吞吐，需要结合实测权衡。

示例四：更进一步，把状态卸载到 NVMe

当 CPU 内存也不够时（典型如数十亿乃至万亿参数模型），只要宿主机配有 NVMe 设备，就可以把优化器状态与参数同时卸载到 NVMe。nvme_path 指向可写的挂载目录：

{
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "nvme",
            "nvme_path": "/nvme_data"
        },
        "offload_param": {
            "device": "nvme",
            "nvme_path": "/nvme_data"
        }
    }
}

卸载到 NVMe 时，可以额外为两个子配置补充 buffer_count、buffer_size、max_in_cpu、fast_init 等参数（默认值见上文表格）。需要理解的是，卸载的本质是用通信/IO 换显存：CPU/NVMe 卸载得越多，单个 step 中参数、梯度的搬运量越大，因此应优先使用通信-计算重叠选项来掩盖延迟。

更进一步：MiCS 配置

MiCS（Minimally 大模型多节点分片，仓库实现位于 deepspeed/runtime/zero/mics.py）的目标是把模型状态的分区范围从整组 GPU 收缩到更小的设备子集，从而降低跨节点通信量。MiCS 假定已启用 ZeRO Stage 3，它复用 DeepSpeedZeroConfig 来承载全部配置，目前只有两个专属字段：

  • mics_shard_size：用于分区模型状态的设备数量。例如设为 8，表示模型状态只在 8 个设备间切分。
  • mics_hierarchical_params_gather：前向计算时是否采用两阶段层次化（hierarchical）的参数收集方式。当模型状态跨多个节点分区、且节点间带宽较慢时该选项很有价值；默认关闭。

典型配置（模型状态分到单节点内的 8 个设备，故层次化收集为 false）：

{
    "zero_optimization": {
        "stage": 3,
        "mics_shard_size": 8,
        "mics_hierarchical_params_gather": false
    }
}

需要指出的是，在 config.py 中 mics_shard_size 的默认值为 -1（即不启用 MiCS 分片语义），mics_hierarchical_params_gather 默认为 false，与文档描述一致。

ZeRO-3 的运行时假设：为何大多数模型无需改动

DeepSpeed 自动以 (子)模块 forward() 方法为粒度协调参数的收集（all-gather）、分区（scatter）与卸载（offload），反向传播同样如此处理。这套机制背后有两个关键假设，理解它们有助于预判何时需要人工介入：

  1. 每个子模块的前向与反向必须能单独放进设备内存。如果某个大线性层本身就无法放入显存，deepspeed.zero.TiledLinear 提供的 memory-centric tiling 可以将线性层拆成一串更小的、能放入内存的连续子块（详见下文）。
  2. 某模块的参数只会在其自身的 __init__ 与 forward() 中被访问。若训练中需要在别处访问参数，就必须让 DeepSpeed 显式收集并再次分区该参数——这正是下一节"手动参数协调"要解决的场景。

构造超大模型：deepspeed.zero.Init 上下文

ZeRO-3 能够承载参数总量超过单个节点内存上限的巨型模型。对于不引入模型并行（model parallelism）的典型训练场景，只需把模型分配放进 deepspeed.zero.Init() 上下文即可：

with deepspeed.zero.Init():
    model = MyLargeModel()

从源码看，deepspeed.zero.Init（deepspeed/runtime/zero/partition_parameters.py，并由 deepspeed/runtime/zero/init.py 导出）是一个上下文管理器：在上下文内实例化的模块参数会被立即分区（partitioned）并以"假分区"状态存在，只有在真正参与前向/反向时才按需收集。因此"构造超大模型"的内存峰值远低于直接 model = MyLargeModel() 的写法。

一个常见的使用姿势是：先 with deepspeed.zero.Init(): 构造模型，再把模型交给 deepspeed.initialize() 得到 ds_engine，之后的训练循环完全不变：

ds_engine, optimizer, _, _ = deepspeed.initialize(
    model=model, model_parameters=model.parameters(), config=ds_config
)

手动参数协调：外部参数的收集与注册

大多数模型无需任何修改即可跑 ZeRO-3，但以下情况需要协调分区权重：

  • 在训练循环之外（如计算指标、日志、可视化、checkpoint 处理）访问模型权重；
  • 训练期间在多个子模块之间共享权重（典型的权重捆绑 / weight tying）。

DeepSpeed 将这类"在被创建它的模块之外被访问的参数"称为 external parameters（外部参数），并提供了若干机制。

收集参数：GatheredParameters

对使用 deepspeed.zero.Init 分区构造的模型，若要访问某模块在其构造函数或 forward() 之外使用的权重，可以借助 GatheredParameters 上下文管理器，把对应参数的完整值在上下文内临时收集（gather）到当前进程：

from deepspeed.zero import GatheredParameters

with GatheredParameters(model.embeddings.weight):
    # 在此范围内，model.embeddings.weight 是可用的完整张量
    do_something_with(model.embeddings.weight)

GatheredParameters 也可用于 register_external_parameter 的补充场景（见下文）。需要注意：上下文管理器内部会发起集合通信，因此需要所有参与训练的进程协同调用。

注册外部参数：register_external_parameter / unregister_external_parameter

ZeRO-3 在前向与反向中按需自动收集、再分区参数；但当某个参数在其所属模块的 forward() 之外被使用，就需要把它注册为外部参数，ZeRO-3 才能正确协调。DeepSpeed 从 0.3.15 版本起支持自动外部参数发现与注册，覆盖了最常见的场景；无法自动识别时仍可手动注册。register_external_parameter 与 unregister_external_parameter 的函数签名如下（实现同样位于 partition_parameters.py）：

  • register_external_parameter(module, parameter)：把参数与需要在其前向/后向期间访问它的模块关联起来；
  • unregister_external_parameter(module, parameter)：解除上述注册。

DeepSpeed 能自动识别的外部参数场景包括：

  1. 参数跨 forward 使用（语言模型中的典型模式）：以 GPT 类模型为例，embeddings.weight 既在 embeddings.forward() 中被使用，又在 compute_logits() 中被使用——后者位于训练循环中、且发生在持有该权重的模块的 forward 之外，因此 embeddings.weight 是一个外部参数：
class LanguageModel(torch.nn.Module):
    ...
    def forward(self, inputs):
        embeds = self.embeddings(inputs)
        ...
        logits = compute_logits(output, self.embeddings.weight)
        ...
  1. forward 中返回参数：下例中 CustomLinear.forward() 把自身 bias 作为返回值之一交给外层使用，DeepSpeed 会检测到这个外部 bias，并将其注册到使用 CustomLinear 的子模块上：
class CustomLinear(torch.nn.Linear):
    def forward(self, *input):
        output = super().forward(*input)
        return output, self.bias

覆盖 Module.apply：override_module_apply

Module.apply(fn) 是 PyTorch 中自定义模型初始化的常用入口。在 ZeRO-3 下，模型初始化期间参数已被 zero.Init 分区，因此 Module.apply 的实现必须顾及这种分区状态。

ZeRO-3 的默认行为是覆盖（override）Module.apply：在参数被 Module.apply 访问之前先自动收集。好处是开发方便——用户不必在 Module.apply 里手工协调分区参数；代价是模型初始化变慢，因为即便常见的用法只是定制少数参数，也会触发全量（数十亿级）参数收集。

若你的初始化代码只触及少量参数，可在配置中把 override_module_apply 设为 false 关闭该默认覆盖，换取更快的模型初始化，代价是必须在自己的 Module.apply 实现中手工处理分区参数。

Memory-Centric Tiling：TiledLinear 与任意尺寸算子

为了降低大模型训练的工作内存（working memory）需求，ZeRO-Infinity 引入了名为 memory-centric tiling 的技术：利用 ZeRO-3 的"数据按需取用、用完即释放"模式，把一个巨大的算子切分成若干可以顺序执行的小 tile，从而把工作内存需求摊薄。与 ZeRO-3 结合后，每个 tile 的参数和梯度可以被"取一块、算一块、释放一块"，工作内存随 tile 数量成比例下降。

因此，ZeRO-Infinity 可以支持任意规模的算子，而无需为了适配 GPU 显存而重构模型做模型并行。落地到代码层面即是 deepspeed.zero.TiledLinear（源码见 deepspeed/runtime/zero/tiling.py）：它实现 memory-centric tiling，与 ZeRO-3 配合把线性层拆解成更小的子模块序列，从而满足前文"子模块前向/反向必须单独放入显存"的假设。

调试 ZeRO：访问完整与分区的模型状态

ZeRO 训练难以调试的根源在于：参数、梯度与优化器状态三者均被分区，任何一类模型状态都无法像普通训练那样直接访问。为此 DeepSpeed 提供了一组例程，允许以**分区（local）与非分区（full）**两种形态访问各类模型状态，其实现位于 deepspeed/utils/tensor_fragment.py。

读取接口

接口语义
safe_get_full_fp32_param(param)组装并返回 fp32 参数完整值
safe_get_full_grad(param)组装并返回完整梯度
safe_get_full_optimizer_state(param, optim_state_key)组装并返回完整 fp32 优化器状态（key 如 Adam 的 exp_avg / exp_avg_sq）
safe_get_local_fp32_param(param)返回本进程持有的 fp32 参数分区（stage 3）
safe_get_local_grad(param)返回本进程的梯度分区（stage 3）
safe_get_local_optimizer_state(param, optim_state_key)返回本进程的优化器状态分区（stage 3）

从源码实现可以精确还原其行为分支：对 ZeRO-3 参数（带 ds_id 属性），走 param._z3_optimizer.get_full_hp_param(...) / get_full_hp_grad(...)；对 Stage 1/2 及 bf16 优化器的参数（带 _hp_mapping 属性），走 param.get_full_hp_param()。若参数既无 ds_id 也无 _hp_mapping，safe_get_full_* 返回 None。特别地，full 梯度通过 all_reduce 汇总各 rank 的分区梯度得到（见 tensor_fragment.py）。

使用上述接口有三个重要注意事项：

  1. 即使对应模型状态被卸载到 CPU 或 NVMe，这些 API 返回的张量也总是位于 accelerator 设备（如 CUDA）上。
  2. 要访问非分区（full）形态，所有参与训练的进程都必须调用这些工具——即使你只在主进程中消费结果；若只有部分进程调用，程序会因等待各方贡献数据而挂起。
  3. 这些例程只在训练的特定阶段返回正确数据：梯度在 backward() 之后、optimizer.step() 之前有效；优化器状态（以及 fp32 主权重）在 step() 之后才完成更新。

训练循环中的读取示例

backward(loss)
[...]
from deepspeed.utils import safe_get_full_fp32_param, safe_get_full_grad, safe_get_full_optimizer_state

for n, lp in model.named_parameters():
    # 1. 访问 full 状态
    #   1.1) 梯度查询：zero1/zero2 需在 backward 后、step 前调用；zero3 在 backward 后即可调用
    hp_grad = safe_get_full_grad(lp)

    #   1.2) fp32 主权重与优化器状态在训练循环中随时可查，但会在 step 之后才更新
    hp = safe_get_full_fp32_param(lp)
    exp_avg = safe_get_full_optimizer_state(lp, "exp_avg")
    exp_avg_sq = safe_get_full_optimizer_state(lp, "exp_avg_sq")

    # 2. 访问 local 状态（zero3）
    #   对于 zero3，参数、梯度、优化器状态全部被分区，各进程只能访问自己那份分区
    local_hp = safe_get_local_fp32_param(lp)
    local_hp_grad = safe_get_local_grad(lp)
    local_exp_avg = safe_get_local_optimizer_state(lp, "exp_avg")
    local_exp_avg_sq = safe_get_local_optimizer_state(lp, "exp_avg_sq")

[...]
optimizer.step()

修改分区状态：梯度裁剪、权重重置等定制场景

有些场景需要在常规训练循环之外修改参数、梯度或优化器状态——例如实现自定义的梯度处理、权重初始化、联邦学习式的聚合等。由于分区机制，这并不直观，为此 DeepSpeed 提供了对应的写接口：

接口语义
safe_set_full_fp32_param(param, value)用 full 值更新 fp32 主参数
safe_set_full_optimizer_state(param, value, optim_state_key)用 full 值更新优化器状态
safe_set_full_grad(param, value)用 full 值更新梯度
safe_set_local_fp32_param(param, value)更新本进程的 fp32 参数分区
safe_set_local_grad(param, value)更新本进程的梯度分区
safe_set_local_optimizer_state(param, value, optim_state_key)更新本进程的优化器状态分区
safe_update_full_grad_vectorized(...)向量化地整体更新 full 梯度

修改参数与优化器状态：initialize() 之后任意时刻可用

参数与优化器状态的修改例程在 deepspeed.initialize() 完成（DeepSpeed 引擎初始化）之后的任何时间点都可用。下面的示例把全部 fp32 参数与 Adam 优化器状态清零：

[...]
from deepspeed.runtime.zero.utils import is_zero_param
from deepspeed.utils import safe_set_full_fp32_param, safe_set_full_optimizer_state
from deepspeed.utils import safe_set_local_fp32_param, safe_set_local_optimizer_state

# 示例：清零所有 fp32 参数与优化器状态
for n, lp in model.named_parameters():
    # 1. 对 zero stage 1/2/3：设置 full fp32 参数与其 full 优化器状态
    zero_tensor = torch.zeros(lp.ds_shape) if is_zero_param(lp) else torch.zeros(lp.shape)

    safe_set_full_fp32_param(lp, zero_tensor)
    safe_set_full_optimizer_state(lp, zero_tensor, "exp_avg")
    safe_set_full_optimizer_state(lp, zero_tensor, "exp_avg_sq")

    # 2. 对 zero stage 3：各进程单独设置自己的 local fp32 参数与 local 优化器状态
    zero_tensor_local = torch.zeros(lp.ds_tensor.shape)

    safe_set_local_fp32_param(lp, zero_tensor_local)
    safe_set_local_optimizer_state(lp, zero_tensor_local, "exp_avg")
    safe_set_local_optimizer_state(lp, zero_tensor_local, "exp_avg_sq")

[...]

is_zero_param(lp) 用于判断参数是否为 ZeRO 分区参数（实现于 deepspeed/runtime/zero/utils.py）：分区参数用 ds_shape（分区逻辑完整形状）构造 full 张量，普通参数直接用 lp.shape。同理，设置 local 分区时用 ds_tensor.shape。

修改梯度：只能在 backward 后、step 前

梯度相关写接口受生命周期约束：必须在 backward(loss) 之后、optimizer.step() 之前调用。以下示例展示如何清零全部梯度：

backward(loss)
[...]
from deepspeed.runtime.zero.utils import is_zero_param
from deepspeed.utils import safe_set_full_grad, safe_set_local_grad

# 示例：清零全部梯度
for n, lp in model.named_parameters():
    # 1. 对 zero stage 1/2/3：设置 full 梯度
    zero_tensor = torch.zeros(lp.ds_shape) if is_zero_param(lp) else torch.zeros(lp.shape)
    safe_set_full_grad(lp, zero_tensor)

    # 2. 对 zero stage 3：各进程设置自己的 local 梯度分区
    zero_tensor_local = torch.zeros_like(lp.ds_tensor.shape)
    safe_set_local_grad(lp, zero_tensor_local)

[...]
optimizer.step()

同样需谨记：full 形态的写入接口需要所有进程共同参与，否则会因集合通信等待而挂起。

GPU 内存管理：empty_partition_cache

默认情况下，ZeRO-3 训练结束时，部分参数可能仍保持"已收集"（未分区）状态而占用 GPU 显存。这是有意为之的优化——如果后续还要恢复训练（resume），这些常驻参数可省去一次全量收集。若想主动清理这些缓存的参数以释放显存，可调用 DeepSpeed 引擎的 empty_partition_cache 方法（实现于 deepspeed/runtime/engine.py）。

with zero.Init():
    model = MyLargeModel()

ds_engine, _, _, _ = deepspeed.initialize(model, ...)
for batch in ...:
    loss = ds_engine(batch)
    ds_engine.backward(batch)
    ds_engine.step()

# 释放模型参数占用的 GPU 缓存内存
ds_engine.empty_partition_cache()

引擎状态的运行时卸载与重载：offload_states / reload_states

除配置驱动的 ZeRO-Infinity 卸载外，DeepSpeed 引擎还提供了一套运行时 API，用于在特定时机把设备（如 CUDA）内存中的一组状态手动搬到其他设备（目前仅支持 CPU 内存），以临时压缩显存占用。其方法签名（实现于 deepspeed/runtime/engine.py）为：

def offload_states(self,
                   include: Container[OffloadStateTypeEnum] = None,
                   device: OffloadDeviceEnum = OffloadDeviceEnum.cpu,
                   pin_memory: bool = True,
                   non_blocking: bool = False) -> None:
    """Offload the engine's states to the specified device.

    Arguments:
        include: Optional. The set of states to offload. If not provided, all states are offloaded.
        device: Optional. The device to move the ZeRO optimizer buffers to. Currently only `OffloadDeviceEnum.cpu` is supported.
        pin_memory: Optional. Whether to pin the memory of the offloaded states.
        non_blocking: Optional. Whether to offload the states asynchronously.
    """

OffloadStateTypeEnum（定义于 deepspeed/runtime/zero/offload_config.py）枚举了可被卸载的状态类别：

  • OffloadStateTypeEnum.optim_states：优化器状态；目前仅支持 DeepSpeed FusedAdam 优化器的状态；
  • OffloadStateTypeEnum.hp_params：fp32 参数；
  • OffloadStateTypeEnum.lp_params：bf16/fp16 参数；
  • OffloadStateTypeEnum.lp_grads：bf16/fp16 梯度；
  • OffloadStateTypeEnum.contiguous_grad_buffer：用于梯度规约的连续梯度缓冲区。

通过 include 参数可以只卸载其中部分状态。需要权衡的是：卸载换来显存，但引入计算/IO 开销。因此该 API 配套了把状态重新载回设备内存的方法：

def reload_states(self, non_blocking: bool = False) -> None:
    """Reload the engine states to the original device.

    Arguments:
        non_blocking: Optional. Whether to offload the states asynchronously.
    """

典型用法是在一次 forward / backward / step 全部完成之后、需要释放显存执行某些大内存操作（如另一次推理、评估、checkpoint 序列化）之前卸载，操作完成后再重载：

# 在 forward、backward 与 step 之后卸载
ds_engine.offload_states(include=[OffloadStateTypeEnum.hp_params, OffloadStateTypeEnum.optim_states])

# 执行某些需要大量设备内存的操作
...

# 把状态重新载回设备内存
ds_engine.reload_states()

如果只关心"当前某类状态分布在哪些设备上"，可用 deepspeed.runtime.zero.offload_states.get_state_devices（实现于 deepspeed/runtime/zero/offload_states.py）：

def get_state_devices(model, state: OffloadStateTypeEnum) -> Set[torch.device]:
    """Retrieve the devices of the specified state of the model.

    Args:
        model (DeepSpeedEngine): The model whose device allocations are to be checked.
        state (OffloadStateTypeEnum): The specific state for which the devices should be retrieved.

    Returns:
        Set[torch.device]: A set of devices of the specified state.
    """

小结：如何为你的场景选择 ZeRO 方案

综合本文内容，可以按如下思路快速决策：

  • 显存可容纳整模型 + 优化器状态：不需要 ZeRO，纯数据并行即可。
  • 需要容纳更大 batch 或更大模型：启用 Stage 1 → Stage 2 分区优化器状态与梯度。
  • 单卡无法容纳模型本身：启用 Stage 3，配合 deepspeed.zero.Init 构造模型；若前向/反向中的某个巨层放不下，考虑 TiledLinear 的 memory-centric tiling。
  • GPU 显存之外还需要进一步压缩 CPU/整体内存：为 offload_optimizer / offload_param 设置 device: "cpu"（ZeRO-Infinity）；仍不够再用 device: "nvme"。
  • 训练循环之外访问或修改权重：区分"只读"（GatheredParameters、safe_get_*）与"写入"（safe_set_*）场景，并牢记生命周期与全进程参与的约束。

无论选择哪种组合，配置都由 zero_optimization 一节驱动、由 DeepSpeedZeroConfig 解析校验，而仓库内 docs/_pages/config-json.md 提供了完整选项清单，是上线前核对参数的最佳参考。

【免费下载链接】DeepSpeed DeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective. 【免费下载链接】DeepSpeed 项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed

Logo
腾讯云开发者社区

腾讯云面向开发者汇聚海量精品云计算使用和开发经验，营造开放的云计算技术生态圈。

更多推荐

Elasticsearch复杂数据类型终极指南：从入门到精通

Elasticsearch作为功能强大的搜索引擎，支持多种复杂数据类型，让开发者能够灵活处理各种结构化和非结构化数据。本文将带你全面了解Elasticsearch中的复杂数据类型，从基础概念到实际应用，助你轻松掌握数据建模的核心技巧。## 内部对象：构建层级化数据结构在Elasticsearch中，对象类型（Object）是最基础的复杂数据类型之一，用于表示具有嵌套关系的数据。例如，我们可

avatar 腾讯云开发者社区

终极指南：Flink SQL连接器版本管理从混乱到有序的升级之路

Apache Flink作为流处理领域的佼佼者，其SQL连接器的版本管理一直是开发者面临的核心挑战。本文将系统讲解Flink SQL连接器版本管理的最佳实践，帮助你轻松应对版本兼容性问题，实现从混乱到有序的升级之旅。## 连接器版本管理的常见痛点 😫在Flink应用开发中，连接器版本管理常常让开发者头疼不已。不同版本的连接器可能导致各种兼容性问题，例如API变更、功能差异甚至运行时错误。

avatar 腾讯云开发者社区

如何快速搭建Neon无服务器PostgreSQL：面向初学者的完整指南

Neon是一款革命性的无服务器PostgreSQL解决方案，它通过分离存储和计算层，实现了自动扩缩容、类代码式数据库分支以及零级扩展能力。本指南将帮助你从零开始搭建Neon开发环境，体验这款创新数据库的强大功能。## 准备工作：环境要求与依赖项在开始搭建Neon环境前，请确保你的系统满足以下要求：- Linux操作系统（推荐Ubuntu 20.04+或Debian 11+）- Git

avatar 腾讯云开发者社区