ZeRO

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

1. ZeRO 第一阶段：优化器状态（例如，对于Adam优化器，32位权重，以及第一和第二矩估计）在进程之间进行分区，以便每个进程仅更新其分区。

2. ZeRO 第二阶段：用于更新模型权重的16位梯度也被分区，使得每个进程仅保留与其优化器状态部分对应的梯度。

3. ZeRO 第三阶段：16位模型参数在进程间进行分区。ZeRO-3 在前向和后向传播过程中会自动收集和分区这些参数。

此外，ZeRO-3 包含了 infinity offload engine 以形成 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

设置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所需的内存

Constraints

• ge = 0

use_multi_rank_bucket_allreduce: bool = True

将不同等级的reduce桶合并，并执行All-Reduce操作，而不是多个Reduce操作。 当模型较小且我们希望在太多GPU上扩展时，此功能非常有用，从而减少每个数据包的消息大小。

allgather_partitions: bool = True

在每一步结束时，从所有GPU中收集更新的参数时，选择使用allgather集体操作或一系列广播集体操作

allgather_bucket_size: int = 500,000,000

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

Constraints

• 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使用

Constraints

• 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。

Constraints

• ge = 0

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

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

Constraints

• ge = 0

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

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

Constraints

• ge = 0

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

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

Constraints

• ge = 0

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

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

Constraints

• 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

针对CPU卸载的第1和第2阶段优化，通过细粒度的梯度分区在多个等级之间并行化梯度复制到CPU内存。性能优势随着梯度累积步骤（在优化器步骤之间更多的复制）或GPU数量（增加的并行性）的增加而增长。

zero_hpz_partition_size: int = 1

零参数分区次级组中的排名数量

Constraints

• 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配置中的示例配置： “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

重写 nn.Module 的 apply 函数，用于阶段 3。

class deepspeed.runtime.zero.config.DeepSpeedZeroOffloadParamConfig

设置参数卸载的选项。仅在阶段3有效。

device: OffloadDeviceEnum = 'none'

用于卸载模型参数的设备内存。支持的选项是cpu和nvme。

nvme_path: Optional[Path] = None

用于参数卸载的NVMe设备的文件系统路径。

buffer_count: int = 5

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

Constraints

• ge = 0

buffer_size: int = 100,000,000

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

Constraints

• ge = 0

max_in_cpu: int = 1,000,000,000

当启用卸载到NVMe时，在CPU内存中维护的参数元素数量。

Constraints

• ge = 0

pin_memory: bool = False

卸载到页面锁定的CPU内存。这可能会以额外的内存开销为代价提高吞吐量。

class deepspeed.runtime.zero.config.DeepSpeedZeroOffloadOptimizerConfig

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

device: OffloadDeviceEnum = 'none'

设备内存用于卸载优化器状态。支持的选项有cpu和nvme。无论选择哪种设备，优化器计算都会卸载到CPU。

nvme_path: Optional[Path] = None

用于优化器状态卸载的NVMe设备的文件系统路径。

buffer_count: int = 4

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

Constraints

• 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的优化器状态百分比。仅在使用ZeRO Stage 3时有效。

Constraints

• ge = 0.0

• le = 1.0

ZeRO-3 配置示例

1. 使用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
           }
       },
       ...
   }
   

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

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

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

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

4. 通过卸载到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配置示例

1. 使用MiCS来分区模型状态（包括优化器状态、梯度和参数）。以下配置示例将模型状态分区到八个设备，并假设这八个设备位于单个节点内（mics_hierarchical_params_gather为False）。

   {
       "zero_optimization": {
           "stage": 3,
           "mics_shard_size": 8,
           "mics_hierarchical_params_gather": False,
       },
       ...
   }
   

假设

DeepSpeed 自动协调参数的收集（即， 全收集）、分区（即， 分散）以及卸载，这些操作在（子）模块的 forward() 方法粒度上进行。反向传播的处理方式类似。这一策略基于两个基本假设：

1. 子模块的前向和后向传递必须分别适应设备内存。如果不是这种情况，deepspeed.zero.TiledLinear实现了以内存为中心的平铺，并与ZeRO-3一起工作，将线性层分解为一系列可以适应内存的较小子模块。

2. 模块的参数仅在其自身的__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 可以自动检测以下外部参数场景：

1. 参数访问：考虑以下在语言模型（如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)
           ...
   

2. 返回一个参数：

   CustomLinear 返回一个输出及其自身的 bias 参数。DeepSpeed 将检测外部的 bias 参数并将其注册到使用 CustomLinear 的子模块中。

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

重写模块.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)

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

Parameters

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

deepspeed.utils.safe_get_full_grad(param)

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

Parameters

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

deepspeed.utils.safe_get_full_optimizer_state(param, optim_state_key)

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

Parameters

• param (torch.nn.Parameter) – 一个模型参数

• optim_state_key (string) – 优化器状态的键值（例如，Adam优化器中的exp_avg）

deepspeed.utils.safe_get_local_fp32_param(param)

获取以fp32精度分区的ZeRO-3分区参数的本地分区。 :param param: 模型参数。 :type param: torch.nn.Parameter

deepspeed.utils.safe_get_local_grad(param)

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

deepspeed.utils.safe_get_local_optimizer_state(param, optim_state_key)

获取以fp32精度分区的ZeRO-3分区参数的本地优化器状态分区。 :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)

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

Parameters

• param (torch.nn.Parameter) – 一个模型参数

• value (torch.Tensor) – 新值

deepspeed.utils.safe_set_full_optimizer_state(param, value, optim_state_key)

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

Parameters

• param (torch.nn.Parameter) – 一个模型参数

• value (torch.Tensor) – 新值

• optim_state_key (string) – 优化器状态的键值（例如，Adam优化器中的exp_avg）

deepspeed.utils.safe_set_full_grad(param, value)

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

Parameters

• param (torch.nn.Parameter) – 一个模型参数

• value (torch.Tensor) – 未分区的新梯度值。

deepspeed.utils.safe_set_local_fp32_param(param, value)

更新ZeRO-3分区参数的本地分区。 :param param: 模型参数。 :type param: torch.nn.Parameter :param value: 本地参数分区的新值。 :type value: torch.Tensor

deepspeed.utils.safe_set_local_grad(param, value)

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

Parameters

• param (torch.nn.Parameter) – 一个模型参数。

• value (torch.Tensor) – 本地梯度分区的新值。

deepspeed.utils.safe_set_local_optimizer_state(param, value, optim_state_key)

更新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.

    """