ZeRO

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

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

2. ZeRO 阶段 2：用于更新模型权重的 32 位梯度的降维也被分区，因此每个进程仅保留与其优化器状态部分相对应的梯度。

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

设置 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 桶，并执行 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 模型权重。

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

零参数分区次级组中的等级数量

约束

• ge = 0

zero_quantized_weights: bool = False

布尔值，指示是否对零参数（权重）进行量化以实现有效的 all_gather 通信

zero_quantized_nontrainable_weights: bool = False

布尔值，指示是否对非训练零参数（权重）进行量化，以实现有效的内存使用和通信。 与 zero_quantized_weights 不同，zero_quantized_weights 会以原始精度存储权重，并且仅在通信期间执行量化，此标志将以量化精度存储权重。 这对于 LoRA 训练非常有用。

zero_quantized_gradients: bool = False

布尔值，指示是否使用量化的零梯度来实现有效的 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

在 Stage 3 中覆盖 nn.Module apply 函数。

class deepspeed.runtime.zero.config.DeepSpeedZeroOffloadParamConfig

为参数卸载设置选项。仅在 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

为优化器卸载设置选项。在 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 配置示例

1. 使用 ZeRO 来划分优化器状态（stage 1）、梯度（stage 2）和参数（stage 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 stage 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()

class deepspeed.zero.Init(module=None, data_parallel_group=None, mem_efficient_linear=True, remote_device=None, pin_memory=False, config_dict_or_path=None, config=None, enabled=True, dtype=None, mpu=None, zero_param_parallel_group=None, zero_quantized_weights=False, zero_quantized_nontrainable_weights=False, sequence_data_parallel_group=None, param_swapper=None)

get_partition_rank()

子类可以重载以指定参数分区组中不同的相对排名

get_dp_process_group()

返回包含所有数据并行排名的通信组

手动参数协调

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

收集参数

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

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

class deepspeed.zero.GatheredParameters(params, modifier_rank=None, fwd_module=None, enabled=True)

注册外部参数

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
   

deepspeed.zero.register_external_parameter(module, parameter)

指示 DeepSpeed 在 module 的正向和反向传递中协调 parameter 的收集和划分。

当参数在其拥有模块的 forward() 之外访问时，使用此方法。DeepSpeed 必须知道何时从其分区状态收集它以及何时释放内存。

注意

这仅适用于使用 ZeRO 阶段 3 进行训练。

参数

• module (torch.nn.Module) – 在其正向传递中需要 parameter 的模块。

• parameter (torch.nn.Parameter) – 要注册的参数。

引发

RuntimeError – 如果 parameter 不是 torch.nn.Parameter 类型。

示例

1. 注册另一个模块的正向传递中使用的权重（第 6 行）。参数 layer1.weight 被 layer2 使用（第 11 行）。

   class ModuleZ3(torch.nn.Module):
       def __init__(self, *args):
           super().__init__(self, *args)
           self.layer1 = SomeLayer()
           self.layer2 = OtherLayer()
           deepspeed.zero.register_external_parameter(self, self.layer1.weight)
   
       def forward(self, input):
           x = self.layer1(input)
           # self.layer1.weight is required by self.layer2.forward
           y = self.layer2(x, self.layer1.weight)
           return y
   

覆盖模块。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 内存中。

class deepspeed.zero.TiledLinear(in_features, out_features, bias=True, in_splits=1, out_splits=1, input_is_already_split=False, combine_out_splits=True, linear_cls=<class 'torch.nn.modules.linear.Linear'>, init_linear=None, **kwargs)

forward(input_)

定义每次调用时执行的计算。

应由所有子类覆盖。

注意

虽然前向传递的配方需要在此函数中定义，但应该随后调用Module实例，而不是此函数，因为前者负责运行已注册的钩子，而后者则静默地忽略它们。

copy_params_from(other)

从other复制权重和偏差数据。

这对于可重复初始化和测试特别有用。

等效于

with torch.no_grad():
    self.weight.copy_(other.weight)
    if self.bias is not None:
        self.bias.copy_(other.bias)

注意

如果启用了 ZeRO-3，这将是一个集体操作，数据并行等级 0 的更新参数将在所有等级上可见。有关更多信息，请参阅 deepspeed.zero.GatheredParameters。

参数

other (torch.nn.Linear) – 要从中复制的线性层。

调试

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

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

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

deepspeed.utils.safe_get_full_fp32_param(param)

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

参数

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

deepspeed.utils.safe_get_full_grad(param)

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

参数

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

deepspeed.utils.safe_get_full_optimizer_state(param, optim_state_key)

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

参数

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

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

deepspeed.utils.safe_get_local_fp32_param(param)

获取 fp32 分区参数。:param param: 模型参数 :type param: torch.nn.Parameter

deepspeed.utils.safe_get_local_grad(param)

获取分区参数的 fp32 梯度。:param param: 模型参数 :type param: torch.nn.Parameter

deepspeed.utils.safe_get_local_optimizer_state(param, optim_state_key)

获取分区参数的 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) 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)


    # 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 参数。

参数

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

• value (torch.Tensor) – 新值

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

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

参数

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

• value (torch.Tensor) – 新值

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

deepspeed.utils.safe_set_local_fp32_param(param, value)

更新分区 fp32 参数。:param param: 模型参数 :type param: torch.nn.Parameter :param value: 新值 :type value: torch.Tensor

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

更新分区参数的 fp32 优化器状态。: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.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 or 2, set the full fp32 and their full optim states
    zero_tensor = torch.zeros_like(lp)

    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_like(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")

[...]

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()