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

每次 reduce/allreduce 的元素数量。限制大型模型 allgather 所需的内存。

约束

ge = 0

use_multi_rank_bucket_allreduce: bool = True: 结合不同 rank 的 reduce 桶并执行 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 副本初始化 fp32 主权重（无精度损失）还是从模型的 fp16 副本初始化（有精度损失）。这可用于在检查点缺少优化器状态时初始化优化器状态。

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 (别名 'stage3_prefetch_bucket_size')

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

约束

ge = 0

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

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

约束

ge = 0

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

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

约束

ge = 0

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

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

约束

ge = 0

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

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

约束

ge = 0

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

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

use_all_reduce_for_fetch_params: bool = False (别名 '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 卸载，通过细粒度梯度分区在 rank 之间并行化梯度复制到 CPU 内存。性能优势随着梯度累积步数（优化器步数之间更多复制）或 GPU 数量（并行度增加）而增长。

zero_hpz_partition_size: int = 1

零参数分区二级组中的 rank 数量。

约束

ge = 0

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

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

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

zeropp_loco_param: Optional[Dict[str, Any]] = None: 此字典包含使用 LoCo-Zero++ 的参数，其中有两个关键参数：- err_beta：梯度计算前后量化误差移动平均的系数。它介于 0 和 1 之间，默认值为 0.8。- reset_T：移动平均误差缓冲区被清除的步数。默认值为 1024。这些参数可以根据性能需求进行调整。ds config 中的配置示例：“zeropp_loco_param”：{“err_beta”：0.8，“reset_T”：1024}。有关更多详细信息，请参阅 LoCo 论文：(https://arxiv.org/abs/2407.04480)。

mics_shard_size: int = -1

mics_hierarchical_params_gather: bool = False

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

pipeline_loading_checkpoint: bool = False

override_module_apply: bool = True: 对于阶段 3，覆盖 nn.Module apply 函数。

log_trace_cache_warnings: bool = False: 是否记录跟踪缓存的警告，例如失效事件。

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

设置参数卸载选项。仅适用于阶段 3。

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

nvme_path: Optional[Path] = None: 用于参数卸载的 NVMe 设备的 filesystem 路径。

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]

设置优化器卸载选项。适用于阶段 1、2 和 3。

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

nvme_path: Optional[Path] = None: 用于优化器状态卸载的 NVMe 设备的 filesystem 路径。

buffer_count: int = 4

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

约束

ge = 0

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

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

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

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

ratio: float = 1.0

卸载到 CPU Adam 的优化器状态百分比。仅适用于 ZeRO 阶段 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() 方法的粒度上自动协调参数的收集（即 all-gather）、分区（即 scatter）和卸载。反向传播也以类似方式处理。此策略有两个基本假设

子模块的前向和反向传播必须单独适应设备内存。如果不是这种情况，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 实现中手动处理分区参数。

以内存为中心的平铺

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

调试

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

重要提示

# 即使相应的模型状态已卸载到 CPU 或 NVMe，这些 API 也会返回加速器设备上的张量。

# 要访问非分区（完整）形式，这些实用程序必须由所有参与训练的进程调用，即使您决定只在主进程中处理结果。如果所有进程不参与，这些实用程序将挂起，等待所有进程发送其贡献。

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

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

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

参数: param (torch.nn.Parameter) – 一个模型参数
返回: 加速器设备上的张量
返回类型: Union[torch.Tensor, None]

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

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

参数: param (torch.nn.Parameter) – 一个模型参数
返回: 加速器设备上的张量
返回类型: Union[torch.Tensor, None]

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）

返回

加速器设备上的张量

返回类型

Union[torch.Tensor, None]

deepspeed.utils.safe_get_local_fp32_param(param)

以 fp32 精度获取 ZeRO-3 分区参数的局部分区。

参数: param (torch.nn.Parameter) – 一个模型参数。
返回: 加速器设备上的张量
返回类型: Union[torch.Tensor, None]

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

获取 ZeRO-3 分区参数的局部梯度分区。返回的数据类型是用于梯度累积的数据类型。这通常是参数数据类型，但也可能不同（例如，bf16 参数训练与 fp32 梯度累积）。

参数: param (torch.nn.Parameter) – 一个模型参数
返回: 加速器设备上的张量
返回类型: Union[torch.Tensor, None]

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

以 fp32 精度获取 ZeRO-3 分区参数的局部优化器状态分区。

参数

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

返回

加速器设备上的张量

返回类型

Union[torch.Tensor, None]

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

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 (torch.nn.Parameter) – 一个模型参数。
value (torch.Tensor) – 局部参数分区的新值。

deepspeed.utils.safe_set_local_grad(param, value)

更新 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 (torch.nn.Parameter) – 一个模型参数。
value (torch.Tensor) – 局部优化器状态分区的新值。
optim_state_key (string) – 优化器状态的键值（例如 Adam 优化器中的 exp_avg）。

deepspeed.utils.safe_update_full_grad_vectorized(param_list: List[Parameter], update_func: Callable)[source]

对低精度（例如 fp16）参数列表的分区梯度进行矢量化更新。为避免精度问题，更新值应具有梯度累积的数据类型。

参数

param_list (List[torch.nn.Parameter]) – 模型参数列表
update_func (torch.Tensor) – 一个函数，它接受当前的完整梯度值并返回新的梯度值。

修改参数和优化器状态的例程可以在 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：用于 reduce 操作的连续梯度缓冲区。

请注意，卸载状态伴随着内存节省和计算开销之间的权衡。此 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.

    """