ZeRO

零冗余优化器 (ZeRO) 通过在数据并行进程之间对三种模型状态（优化器状态、梯度和参数）进行分区，而不是复制它们，从而消除了数据并行进程之间的内存冗余。通过这样做，与经典的数据并行相比，它提高了内存效率，同时保留了其计算粒度和通信效率。

**ZeRO 阶段 1**：优化器状态（例如，对于Adam 优化器，32 位权重以及第一和第二矩估计）在进程之间进行分区，以便每个进程仅更新其分区。
**ZeRO 阶段 2**：用于更新模型权重的缩减的 16 位梯度也进行分区，以便每个进程仅保留与其优化器状态部分相对应的梯度。
**ZeRO 阶段 3**：16 位模型参数在进程之间进行分区。ZeRO-3 将在正向和反向传递期间自动收集和分区它们。

此外，ZeRO-3 包括 *无限卸载引擎* 来形成 ZeRO-Infinity（[论文](https://arxiv.org/abs/2104.07857)），它可以将所有模型状态卸载到 CPU 和 NVMe 内存中，从而节省大量内存。

有关我们算法的深入探讨，请参阅我们关于ZeRO、ZeRO-Offload 和ZeRO-Infinity 的论文。

注意

DeepSpeed 最初在 **ZeRO-Offload** 中包含了卸载功能，这是一个用于在 ZeRO-2 中将优化器和梯度状态卸载到 CPU 内存的系统。**ZeRO-Infinity** 是下一代卸载功能，可用于 ZeRO-3。ZeRO-Infinity 拥有 ZeRO-Offload 的所有节省功能，此外还能够卸载更多模型权重，并具有更有效的带宽利用率和计算与通信的重叠。

入门

如果您不熟悉 DeepSpeed，请查看我们的入门页面。

在使用 DeepSpeed 进行训练后，启用 ZeRO-3 卸载就像在您的 DeepSpeed 配置中启用它一样简单！下面是几个 ZeRO-3 配置示例。请参阅我们的配置指南，以获取配置和性能调整的完整选项列表。

注意

ZeRO-Infinity 和 ZeRO-Offload 与我们高度优化的deepspeed.ops.adam.DeepSpeedCPUAdam 优化器配合使用效果最佳。我们建议使用我们的优化器配置来指示 deepspeed.initialize() 为您构建优化器。

ZeRO 配置

DeepSpeed ZeRO 的所有设置都使用DeepSpeedZeroConfig 进行设置。在主 DeepSpeed 配置字典的 zero_optimization 条目下提供的字典将使用此类进行解析和验证。参数卸载和优化器卸载设置的子配置由DeepSpeedZeroOffloadParamConfig 和DeepSpeedZeroOffloadOptimizerConfig 进行解析。

class deepspeed.runtime.zero.config.DeepSpeedZeroConfig[source]

设置 ZeRO 优化的参数。

stage: ZeroStageEnum = 0: 选择 ZeRO 优化器的不同阶段。阶段 0、1、2 和 3 分别表示禁用、优化器状态分区、优化器+梯度状态分区以及优化器+梯度+参数分区。

contiguous_gradients: bool = True: 在生成梯度时将其复制到连续缓冲区。避免在反向传播过程中出现内存碎片。

reduce_scatter: bool = True: 使用 reduce 或 reduce scatter 而不是 allreduce 来平均梯度

reduce_bucket_size: int = 500,000,000

一次减少/全部减少的元素数量。限制大型模型尺寸的 allgather所需的内存

约束

ge = 0

use_multi_rank_bucket_allreduce: bool = True: 组合不同等级的 reduce bucket 并进行 All-Reduce，而不是多个 Reduce 操作。此功能在模型较小且我们希望将其扩展到太多 GPU 上时很有用，因此会减少每个数据包的消息大小。

allgather_partitions: bool = True: 在 allgather 集合或一系列广播集合之间进行选择，以在每个步骤结束时从所有 GPU 收集更新的参数

allgather_bucket_size: int = 500,000,000

一次 allgather 的元素数量。限制大型模型尺寸的 allgather所需的内存

约束

ge = 0

overlap_comm: Optional[bool] = None: 尝试将梯度的缩减与反向计算重叠

load_from_fp32_weights: bool = True: 布尔值，指示是否从检查点中的 fp32 副本（无精度损失）或模型的 fp16 副本（有精度损失）初始化 fp32 主权重。这可用于即使检查点缺少优化器状态时也能初始化优化器状态。

elastic_checkpoint: bool = False: 启用加载由具有不同 GPU 数量的作业保存的检查点。不再支持。

offload_param: Optional[DeepSpeedZeroOffloadParamConfig] = None: 启用将模型参数卸载到 CPU 或 NVMe。这可以释放 GPU 内存，以便使用更大的模型或批次大小。仅在阶段 3 中有效。预期包含以下值的字典 DeepSpeedZeroOffloadParamConfig。

offload_optimizer: Optional[DeepSpeedZeroOffloadOptimizerConfig] = None: 启用将优化器状态卸载到 CPU 或 NVMe，并将优化器计算卸载到 CPU。这可以释放 GPU 内存，以便使用更大的模型或批次大小。适用于 ZeRO 阶段 1、2、3。预期包含以下值的字典 DeepSpeedZeroOffloadOptimizerConfig。

sub_group_size: int = 1,000,000,000

用于处理参数以适应大型模型（具有数万亿个参数）的切片大小。由 ZeRO3-Offload 和 ZeRO-Infinity 使用。

约束

ge = 0

cpu_offload_param: Optional[bool] = None: 已弃用，请使用 offload_param

cpu_offload_use_pin_memory: Optional[bool] = None: 已弃用，请使用 offload_param 或 offload_optimizer

cpu_offload: Optional[bool] = None: 已弃用，请使用 offload_optimizer

prefetch_bucket_size: int = 50,000,000 (alias 'stage3_prefetch_bucket_size')

在使用前预取的参数元素的最大数量。由 ZeRO3、ZeRO3-Offload、ZeRO-Infinity 和 ZeRO-Inference 使用。

约束

ge = 0

param_persistence_threshold: int = 100,000 (alias 'stage3_param_persistence_threshold')

不要对小于此阈值的参数进行分区。较小的值使用更少的内存，但会大大增加通信量（尤其是延迟受限的消息）。

约束

ge = 0

model_persistence_threshold: int = sys.maxsize (alias 'stage3_model_persistence_threshold')

可以在 GPU 中持久化且不进行分区的参数元素的最大数量。这为由 param_persistence_threshold 设置产生的未分区参数的数量施加了上限。由 ZeRO3-Offload、ZeRO-Infinity 和 ZeRO-Inference 使用。

约束

ge = 0

max_live_parameters: int = 1,000,000,000 (alias 'stage3_max_live_parameters')

在释放之前每个 GPU 上驻留的参数最大数量。较小的值使用更少的内存，但执行更多的通信。

约束

ge = 0

max_reuse_distance: int = 1,000,000,000 (alias 'stage3_max_reuse_distance')

如果参数将在此阈值内的参数内重用，则不要释放它。较小的值使用更少的内存，但执行更多的通信。

约束

ge = 0

gather_16bit_weights_on_model_save: bool = False (alias 'stage3_gather_16bit_weights_on_model_save'): 通过 save_16bit_model() 保存模型之前整合权重。由于权重在 GPU 上进行分区，因此它们不是 state_dict 的一部分，因此当启用此选项时，此函数会自动收集权重，然后保存 fp16 模型权重。

module_granularity_threshold: int = 0 (alias 'stage3_module_granularity_threshold'): 模块的粒度由“参数数量 /（1 + 后代数量）”的比率决定。ZeRO3 将粒度低于阈值的模块分类为细粒度，在获取参数期间将这些模块视为整体单元。这减少了主机开销以及在获取参数时挂钩引入的细粒度层的单独 allgather 开销。

use_all_reduce_for_fetch_params: bool = False (alias 'stage3_use_all_reduce_for_fetch_params'): 在阶段 3 获取模块参数时使用 all_reduce 操作。这可以通过减少主机上连接和切片的开销来提高性能。

stage3_gather_fp16_weights_on_model_save: bool = False: 已弃用，请使用 gather_16bit_weights_on_model_save

ignore_unused_parameters: bool = True: 模块中未使用的参数在静态网络中可能是意外的，但在动态网络中可能是正常的。这控制了当检测到未使用的参数时训练是否应以错误消息终止。默认情况下设置为 True，这意味着忽略未使用的参数并继续训练。现在仅在阶段 2 中使用。

legacy_stage1: bool = False: 为了向后兼容，启用旧的 ZeRO 阶段 1 实现。自行承担风险使用，很快就会弃用。

round_robin_gradients: bool = False: 阶段 1 和 2 优化用于 CPU 卸载，通过细粒度梯度分区将梯度复制到 CPU 内存并行化。性能优势随着梯度累积步数（优化器步骤之间更多的复制）或 GPU 数量（增加并行性）而增长。

zero_hpz_partition_size: int = 1

zero 参数分区辅助组中的排名数量

约束

ge = 0

zero_quantized_weights: bool = False: 布尔值，指示是否对 Zero 参数（权重）进行量化，以实现高效的 all_gather 通信。

zero_quantized_nontrainable_weights: bool = False: 布尔值，指示是否对不可训练的 Zero 参数（权重）进行量化，以实现高效的内存使用和通信。与 `zero_quantized_weights` 不同，后者以原始精度存储权重，仅在通信期间执行量化，而此标志将以量化精度存储权重。这对于 LoRA 训练很有用。

zero_quantized_gradients: bool = False: 布尔值，指示是否使用量化的 Zero 梯度以实现高效的 all_2_all_reduce 通信。

mics_shard_size: int = -1

mics_hierarchical_params_gather: bool = False

memory_efficient_linear: bool = True: 使用内存高效的线性实现，适用于 Stage 3。

pipeline_loading_checkpoint: bool = False

override_module_apply: bool = True: 覆盖 `nn.Module` 的 `apply` 函数，适用于 Stage 3。

class deepspeed.runtime.zero.config.DeepSpeedZeroOffloadParamConfig[source]

设置参数卸载选项。仅在 Stage 3 中有效。

device: OffloadDeviceEnum = 'none': 卸载模型参数的目标设备内存。支持的选项包括 `cpu` 和 `nvme`。

nvme_path: Optional[Path] = None: 参数卸载到 NVMe 设备的文件系统路径。

buffer_count: int = 5

参数卸载到 NVMe 的缓冲池中的缓冲区数量。

约束

ge = 0

buffer_size: int = 100,000,000

参数卸载到 NVMe 的缓冲池中缓冲区的大小。

约束

ge = 0

max_in_cpu: int = 1,000,000,000

启用 NVMe 卸载时，在 CPU 内存中保留的参数元素数量。

约束

ge = 0

pin_memory: bool = False: 卸载到页面锁定的 CPU 内存。这可以提高吞吐量，但会增加额外的内存开销。

class deepspeed.runtime.zero.config.DeepSpeedZeroOffloadOptimizerConfig[source]

设置优化器卸载选项。在 Stage 1、2 和 3 中有效。

device: OffloadDeviceEnum = 'none': 卸载优化器状态的目标设备内存。支持的选项包括 `cpu` 和 `nvme`。无论设备选项如何，优化器计算都会卸载到 CPU。

nvme_path: Optional[Path] = None: 优化器状态卸载到 NVMe 设备的文件系统路径。

buffer_count: int = 4

优化器状态卸载到 NVMe 的缓冲池中的缓冲区数量。这至少应为优化器为每个参数维护的状态数量。例如，Adam 优化器有 4 个状态（参数、梯度、动量和方差）。

约束

ge = 0

pin_memory: bool = False: 卸载到页面锁定的 CPU 内存。这可以提高吞吐量，但会增加额外的内存开销。

pipeline_read: bool = False: 对于基于平铺的优化器步骤处理，将下一个平铺的读取与当前平铺的计算重叠。用于 ZeRO-Infinity。

pipeline_write: bool = False: 对于基于平铺的优化器步骤处理，将上一个平铺的写入与当前平铺的计算重叠。

fast_init: bool = False: 启用 NVMe 卸载时的快速优化器初始化。

ratio: float = 1.0

卸载到 CPU Adam 的优化器状态的百分比。仅在 ZeRO Stage 3 中有效。

约束

ge = 0.0
le = 1.0

ZeRO-3 配置示例

使用 ZeRO 将优化器状态（阶段 1）、梯度（阶段 2）和参数（阶段 3）进行分区。

{
    "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
        }
    },
    ...
}

此外，使用 ZeRO-Infinity 将优化器状态和计算卸载到 CPU。

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

通过将参数卸载到 CPU 内存，节省更多内存。

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

如果您的系统上可用，通过卸载到 NVMe 可以节省更多内存。

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

MiCS 配置

所有 MiCS 配置都使用 DeepSpeedZeroConfig 设置。MiCS 假设启用了 ZeRO 阶段 3 优化。目前，MiCS 有两个配置字段：mics_shard_size 和 mics_hierarchical_params_gather。mics_shard_size 控制用于对模型状态进行分区的设备数量。mics_hierarchical_params_gather 控制我们在前向计算中是否使用两阶段分层方式来收集参数。mics_hierarchical_params_gather 在模型状态跨多个节点分区且跨节点带宽较慢时很有用。默认情况下，此选项处于关闭状态。

MiCS 配置示例

使用 MiCS 对模型状态（包括优化器状态、梯度和参数）进行分区。以下配置示例将模型状态分区到八个设备，并假设这八个设备位于单个节点内（mics_hierarchical_params_gather 为 False）。
{ "zero_optimization": { "stage": 3, "mics_shard_size": 8, "mics_hierarchical_params_gather": False, }, ... }

假设

DeepSpeed 自动协调参数的收集（即全局聚合）、分区（即散射）和卸载，其粒度为（子）模块的 forward() 方法。反向传递以类似的方式处理。此策略有两个基本假设

子模块的前向和反向传递必须分别适合设备内存。如果不是这种情况，deepspeed.zero.TiledLinear 实现以内存为中心的平铺并与 ZeRO-3 一起工作，将线性层分解成一系列可以放入内存的较小子模块。
模块的参数仅在其自己的 __init__ 和 forward() 方法中访问。否则，必须指示 DeepSpeed 收集并重新分区参数。有关手动协调参数的信息，请参阅手动参数协调。

构建大型模型

ZeRO-3 使得参数大小超过系统中单个节点的大型模型成为可能。对于没有模型并行的典型训练情况，您只需在我们的上下文中分配您的模型即可

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

手动参数协调

大多数模型无需修改即可使用 ZeRO-3 进行训练。但是，在某些情况下，可能需要在训练循环之外访问模型权重，或者在训练期间跨子模块共享权重。DeepSpeed 有几种机制可以协调 ZeRO-3 的分区权重。

收集参数

DeepSpeed 提供了收集（或聚合）分区参数的机制。

某些使用 deepspeed.zero.Init 分区的模型可能需要在类构造函数或其 forward() 方法之外访问模块的权重。我们将这些权重称为外部参数，因为这些参数是在创建它们的模块之外访问的。为此，请使用 deepspeed.zero.GatheredParameters 或 deepspeed.zero.register_external_parameter()。

注册外部参数

ZeRO-3 会在模型的前向和反向传递过程中根据需要自动收集和分区模型参数。但是，在某些情况下，参数可能在其模块的前向传递之外使用。我们将这些参数称为外部参数。如果这些参数是自动或手动注册的，ZeRO-3 可以协调这些参数。

注意

DeepSpeed 版本 0.3.15 包括自动外部参数发现和注册，以支持最常见的情况。如果无法自动检测参数，仍然可以手动注册参数。

DeepSpeed 可以自动检测以下外部参数场景

参数访问：考虑以下在 GPT 等语言模型中常见的模式

张量 embeddings.weight 在 embeddings.forward() 和 compute_logits() 中都使用。我们将 embeddings.weight 称为外部参数，因为它是在其拥有模块的前向传递之外的训练循环中使用的。
```
class LanguageModel(torch.nn.Module):
    ...
    def forward(self, inputs):
        embeds = self.embeddings(inputs)
        ...
        logits = compute_logits(output, self.embeddings.weight)
        ...
```

返回参数

CustomLinear 返回输出及其自己的 bias 参数。DeepSpeed 将检测外部 bias 参数并将其与使用 CustomLinear 的子模块注册。

class CustomLinear(torch.nn.Linear):
    def forward(self, *input):
        output = super().forward(*input)
        return output, self.bias

覆盖 Module.apply

自定义模型初始化的一种便捷机制是Module.apply。使用 ZeRO 阶段 3 时，Module.apply 实现必须在模型初始化期间考虑 zero.Init 对参数进行的分区。ZeRO 阶段 3 的默认行为是在覆盖 Module.apply 以确保在 Module.apply 访问之前收集参数，从而自动处理此问题。这种方法的好处是开发方便，因为用户无需在 Module.apply 中手动协调参数。但是，缺点是模型初始化速度慢，因为即使 Module.apply 的常用用法是自定义一些参数，也会收集所有模型参数（例如，数十亿个）。开发人员可以通过将 override_module_apply 配置旋钮设置为 False 来禁用此默认行为，从而加快模型初始化速度，但代价是在其 Module.apply 实现中手动处理分区参数。

以内存为中心的平铺

为了减少大型模型深度学习训练的工作内存需求，ZeRO-Infinity 包含一种称为以内存为中心的平铺的技术，该技术利用 ZeRO-3 的数据获取和释放模式，通过将大型运算符分解成可以顺序执行的较小块来减少工作内存需求。当与 ZeRO-3 结合使用时，每个块的参数和梯度可以一次一个地获取和释放，从而减少与块数成比例的工作内存。因此，ZeRO-Infinity 可以支持任意大小的运算符，而无需重构模型并行性以使其适应有限的 GPU 内存。

调试

参数、梯度和优化器状态的分区使 ZeRO 训练的调试变得复杂。由于此原因，这三组张量（模型状态）都不能正常访问。为了克服这个问题，DeepSpeed 提供了以下例程，用于以其分区（本地）和未分区（完整）形式访问各个模型状态。

重要：请注意，要访问未分区（完整）形式，这些实用程序必须由参与训练的所有进程调用，即使您决定仅在主进程中对结果执行某些操作也是如此。如果所有进程都不参与，这些实用程序将挂起，等待所有进程发送其贡献。

此外，您必须意识到，这些例程仅在训练的特定阶段返回正确的数据。例如，梯度在 backward 之后和 step 之前有效。优化器状态在 step 之后更新。fp32 主权重也是如此。

deepspeed.utils.safe_get_full_fp32_param(param)[source]

组装并返回低精度（例如，fp16）参数的 fp32 参数。

参数: param (torch.nn.Parameter) – 模型参数

deepspeed.utils.safe_get_full_grad(param)[source]

组装并返回低精度（例如，fp16）参数的 fp32 梯度。返回的数据类型是用于梯度累积的类型。这通常是参数数据类型，但也可能不同（例如，使用 fp32 梯度累积的 bf16 参数训练）。

参数: param (torch.nn.Parameter) – 模型参数

deepspeed.utils.safe_get_full_optimizer_state(param, optim_state_key)[source]

组装并返回低精度（例如，fp16）参数的 fp32 优化器状态。

参数

param (torch.nn.Parameter) – 模型参数
optim_state_key (string) – 优化器状态的键值（例如，Adam 优化器中的 exp_avg）

deepspeed.utils.safe_get_local_fp32_param(param)[source]: 获取 ZeRO-3 分区参数的本地分区，以 fp32 精度表示。:param param: 模型参数。:type param: torch.nn.Parameter

deepspeed.utils.safe_get_local_grad(param)[source]: 获取 ZeRO-3 分区参数的本地梯度分区。返回的数据类型是用于梯度累积的类型。这通常是参数数据类型，但也可能不同（例如，使用 fp32 梯度累积的 bf16 参数训练）。:param param: 模型参数 :type param: torch.nn.Parameter

deepspeed.utils.safe_get_local_optimizer_state(param, optim_state_key)[source]: 获取 ZeRO-3 分区参数的本地优化器状态分区，以 fp32 精度表示。:param param: 模型参数 :type param: torch.nn.Parameter :param optim_state_key: 优化器状态的键值（例如，Adam 优化器中的 exp_avg） :type optim_state_key: string

这些例程可以在训练循环中使用，如下面的代码片段所示。

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. Access the full states
    #  1.1) gradient lookup
    # For zero1 and zero2, gradient lookup must be called after `backward` and before `step`
    # For zero3, gradient lookup must be called after `backward`
    hp_grad = safe_get_full_grad(lp)


    # 1.2) fp32 and optim states can probably be called anywhere in the training loop, but will be updated after `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. Access the local states (zero3)
    # For zero3, all of the parameters, gradients, and optimizer states are partitioned,
    # and each process can access its corresponding local state.
    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()

修改分区状态

有时，用户可能希望在常规训练循环之外修改参数、梯度或优化器状态。由于分区，这在 ZeRO 训练中目前比较困难。为了克服这个问题，DeepSpeed 提供了以下例程来修改 fp32 主参数和 fp32 优化器状态。

deepspeed.utils.safe_set_full_fp32_param(param, value)[source]

更新低精度（例如，fp16）参数的分区 fp32 参数。

参数

param (torch.nn.Parameter) – 模型参数
value (torch.Tensor) – 新值

deepspeed.utils.safe_set_full_optimizer_state(param, value, optim_state_key)[source]

更新低精度（例如，fp16）参数的分区 fp32 优化器状态。

参数

param (torch.nn.Parameter) – 模型参数
value (torch.Tensor) – 新值
optim_state_key (string) – 优化器状态的键值（例如，Adam 优化器中的 exp_avg）

deepspeed.utils.safe_set_full_grad(param, value)[source]

更新低精度（例如，fp16）参数的分区梯度。为了避免精度问题，更新值应具有梯度累积的数据类型。

参数

param (torch.nn.Parameter) – 模型参数
value (torch.Tensor) – 未分区的新的梯度值。

deepspeed.utils.safe_set_local_fp32_param(param, value)[source]: 更新 ZeRO-3 分区参数的本地分区。:param param: 模型参数。 :type param: torch.nn.Parameter :param value: 本地参数分区的新的值。 :type value: torch.Tensor

deepspeed.utils.safe_set_local_grad(param, value)[source]

更新 ZeRO-3 分区参数的本地梯度分区。为了避免精度问题，更新值应具有梯度累积的数据类型。

参数

param (torch.nn.Parameter) – 模型参数。
value (torch.Tensor) – 本地梯度分区的新的值。

deepspeed.utils.safe_set_local_optimizer_state(param, value, optim_state_key)[source]: 更新 ZeRO-3 分区参数的本地优化器状态分区。:param param: 模型参数。 :type param: torch.nn.Parameter :param value: 本地优化器状态分区的新的值。 :type value: torch.Tensor :param optim_state_key: 优化器状态的键值（例如，Adam 优化器中的 exp_avg）。 :type optim_state_key: string

修改参数和优化器状态的例程可以在 DeepSpeed 引擎初始化（即 deepspeed.initialize()）之后任何时候使用，如下面的代码片段所示。

[...]
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
# Here is an example to zero all the fp32 parameters and optimizer states.
for n, lp in model.named_parameters():
    # 1. For zero stage 1, 2, or 3 set the full fp32 and their full optim states
    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_get_full_optimizer_state(lp, zero_tensor, "exp_avg")
    safe_get_full_optimizer_state(lp, zero_tensor, "exp_avg_sq")

    # 2. For zero stage 3, each process sets its local fp32 parameters and their local optimizer states individually
    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")

[...]

修改梯度的例程可以在 backward 之后但 step 之前使用，如下面的代码片段所示。

backward(loss)
[...]
from deepspeed.runtime.zero.utils import is_zero_param
from deepspeed.utils import safe_set_full_grad, safe_set_local_grad
# Here is an example of how to zero all the gradients.
for n, lp in model.named_parameters():
    # 1. For zero stage 1, 2, or 3 set the full gradient.
    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. For zero stage 3, each process sets its local gradient partition.
    zero_tensor_local = torch.zeros_like(lp.ds_tensor.shape)

    safe_set_local_grad(lp, zero_tensor_local)

[...]
optimizer.step()

GPU 内存管理

默认情况下，在使用 ZeRO 第 3 阶段完成训练后，某些参数可能保持未分区状态并占用一些 GPU 内存。这样做是为了优化，以防您再次恢复训练。如果您想清除占用 GPU 内存的缓存参数，可以调用 DeepSpeed 引擎的 empty_partition_cache 方法。

以下代码片段说明了此功能。

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

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

# Free GPU memory consumed by model parameters
ds_engine.empty_partition_cache()

卸载状态

DeepSpeed 引擎在设备内存（例如，CUDA 内存）中维护一组状态。以下 API 允许您将这些状态卸载到其他设备（目前，仅支持 CPU 内存），从而减少设备上的内存占用。

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.
    """

您可以通过在 include 参数中指定 OffloadStateTypeEnum 来选择性地卸载特定状态。 OffloadStateTypeEnum 是一个枚举，定义了可以卸载的状态。支持以下状态

OffloadStateTypeEnum.optim_states：优化器状态。目前，仅支持 DeepSpeed 的 FusedAdam 优化器的状态。
OffloadStateTypeEnum.hp_params：FP32 参数。
OffloadStateTypeEnum.lp_params：BF16/FP16 参数。
OffloadStateTypeEnum.lp_grads：BF16/FP16 梯度。
OffloadStateTypeEnum.contiguous_grad_buffer：用于减少操作的连续梯度缓冲区。

请注意，卸载状态需要在内存节省和计算开销之间进行权衡。此 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.
    """

以下是一个示例代码片段，演示如何将 FP32 参数和优化器状态卸载到 CPU 内存

# Offload after forward, backward, and step
ds_engine.offload_states(include=[OffloadStateTypeEnum.hp_params, OffloadStateTypeEnum.optim_states])

# Do something requiring a lot of device memory
...
# Load states back to device memory
ds_engine.reload_states()

deepspeed.runtime.zero.offload_states.get_state_devices 返回指定状态的设备。

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.

    """