零冗余优化器

如果您尚未阅读 DeepSpeed 关于入门和Megatron-LM GPT-2的教程，我们建议您先阅读这些教程，然后再进行本教程。

在本教程中，我们将把 ZeRO 优化器应用于 Megatron-LM GPT-2 模型。ZeRO 是一套强大的内存优化技术，可以有效训练拥有数万亿参数的大型模型，例如 GPT-2 和 Turing-NLG 17B。与训练大型模型的替代模型并行方法相比，ZeRO 的一个关键吸引力在于它不需要修改模型代码。正如本教程将演示的，在 DeepSpeed 模型中使用 ZeRO 既快速又简单，因为您只需更改 DeepSpeed 配置 JSON 中的几个配置。不需要代码更改。

ZeRO 概述

ZeRO 利用数据并行的聚合计算和内存资源，以减少用于模型训练的每个设备（GPU）的内存和计算需求。ZeRO 通过将各种模型训练状态（权重、梯度和优化器状态）分区到分布式训练硬件中可用的设备（GPU 和 CPU）上，从而减少每个 GPU 的内存消耗。具体来说，ZeRO 是作为优化增量阶段实现的，其中早期阶段的优化在后期阶段也可用。要深入了解 ZeRO，请参阅我们的论文。

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

  • 阶段 2：用于更新模型权重的 16 位归约梯度也进行分区，使得每个进程只保留与其部分优化器状态相对应的梯度。

  • 阶段 3：16 位模型参数在进程间进行分区。ZeRO-3 将在正向和反向传播过程中自动收集和分区它们。

此外，ZeRO-3 还包括无限卸载引擎，以形成 ZeRO-Infinity（论文），它可以卸载到 CPU 和 NVMe 内存，从而节省大量内存。

训练环境

我们使用 DeepSpeed Megatron-LM GPT-2 代码进行此练习。您可以参考 Megatron-LM 教程来熟悉代码。我们将在本教程中使用配备 32GB RAM 的 NVIDIA Tesla V100-SXM3 Tensor Core GPU 训练模型。

启用 ZeRO 优化

要为 DeepSpeed 模型启用 ZeRO 优化，我们只需在 DeepSpeed JSON 配置中添加 zero_optimization 键。zero_optimization 键的配置参数的完整说明可在此处获取。

训练一个 1.5B 参数的 GPT-2 模型

我们通过展示 ZeRO 阶段 1 如何在八个 V100 GPU 上实现 15 亿参数 GPT-2 模型的数据并行训练，来证明其优势。我们将训练配置为每个设备使用批大小为 1，以确保内存消耗主要由模型参数和优化器状态引起。我们通过对 DeepSpeed 启动脚本进行以下修改来创建此训练场景：

       --model-parallel-size 1 \
       --num-layers 48 \
       --hidden-size 1600 \
       --num-attention-heads 16 \
       --batch-size 1 \
       --deepspeed_config ds_zero_stage_1.config \

不使用 ZeRO 训练此模型会因内存不足 (OOM) 错误而失败，如下所示：

该模型无法适应 GPU 内存的一个主要原因是，模型的 Adam 优化器状态消耗了 18GB 内存；这占 32GB RAM 的很大一部分。通过使用 ZeRO 阶段 1 在八个数据并行秩之间分区优化器状态，每个设备的内存消耗可以减少到 2.25GB，从而使模型可训练。要启用 ZeRO 阶段 1，我们只需更新 DeepSpeed JSON 配置文件，如下所示：

{
    "zero_optimization": {
        "stage": 1,
        "reduce_bucket_size": 5e8
    }
}

如上所示，我们在 zero_optimization 键中设置了两个字段。具体来说，我们将 stage 字段设置为 1，并将可选的用于梯度归约的 reduce_bucket_size 设置为 500M。启用 ZeRO 阶段 1 后，模型现在可以在 8 个 GPU 上顺利训练而不会耗尽内存。下面我们提供一些模型训练的屏幕截图：

从上面的 nvidia-smi 屏幕截图可以看出，只有 GPU 6-7 用于训练模型。通过 ZeRO 阶段 1，我们可以通过增加数据并行度来进一步减少每个设备的内存消耗。这些节省的内存可以用于增加模型大小和/或批处理大小。相比之下，仅凭数据并行无法实现此类好处。

训练一个 10B 参数的 GPT-2 模型

ZeRO 阶段 2 优化进一步增加了可以使用数据并行训练的模型大小。我们通过使用 32 个 V100 GPU 训练一个包含 10B 参数的模型来展示这一点。

首先，我们需要配置一个启用激活检查点的 10B 参数模型。这可以通过对 DeepSpeed 启动脚本应用以下 GPT-2 模型配置更改来完成：

       --model-parallel-size 1 \
       --num-layers 50 \
       --hidden-size 4096 \
       --num-attention-heads 32 \
       --batch-size 1 \
       --deepspeed_config ds_zero_stage_2.config \
       --checkpoint-activations

接下来，我们需要更新 DeepSpeed JSON 配置，如下所示，以启用 ZeRO 阶段 2 优化：

{
    "zero_optimization": {
        "stage": 2,
        "contiguous_gradients": true,
        "overlap_comm": true,
        "reduce_scatter": true,
        "reduce_bucket_size": 5e8,
        "allgather_bucket_size": 5e8
    }
}

在上述更改中，我们将 stage 字段设置为 2，并配置了 ZeRO 阶段 2 中可用的其他优化参数。例如，我们启用了 contiguous_gradients 以减少反向传播过程中的内存碎片。这些优化参数的完整说明可在此处获取。通过这些更改，我们现在可以启动训练运行。

以下是训练日志的屏幕截图：

以下是显示训练期间 GPU 活动的 nvidia-smi 屏幕截图：

使用 ZeRO-Infinity 训练万亿级模型

ZeRO 的第三阶段 ZeRO-3 将完整的模型状态（即权重、梯度和优化器状态）分区，以使内存节省与数据并行度呈线性关系。ZeRO-3 可以在 JSON 配置中启用。这些配置的完整说明可在此处获取。

使用 ZeRO-Infinity 卸载到 CPU 和 NVMe

ZeRO-Infinity 使用 DeepSpeed 的无限卸载引擎将完整的模型状态卸载到 CPU 或 NVMe 内存，从而支持更大的模型大小。卸载可以在 DeepSpeed 配置中启用：

{
    "zero_optimization": {
        "stage": 3,
        "contiguous_gradients": true,
        "stage3_max_live_parameters": 1e9,
        "stage3_max_reuse_distance": 1e9,
        "stage3_prefetch_bucket_size": 1e7,
        "stage3_param_persistence_threshold": 1e5,
        "reduce_bucket_size": 1e7,
        "sub_group_size": 1e9,
        "offload_optimizer": {
            "device": "cpu"
         },
        "offload_param": {
            "device": "cpu"
       }
   }
}

ZeRO-Infinity 与 ZeRO-Offload： DeepSpeed 最初在 ZeRO-2 中包含了 ZeRO-Offload 卸载功能，这是一个将优化器和梯度状态卸载到 CPU 内存的系统。ZeRO-Infinity 是 ZeRO-3 可访问的下一代卸载功能。ZeRO-Infinity 能够比 ZeRO-Offload 卸载更多的数据，并具有更有效的带宽利用率以及计算和通信的重叠。

分配大规模 Megatron-LM 模型

我们对模型初始化进行了另外两项更改，以支持超出本地系统内存但未超出总系统内存的模型。

  1. 以内存可扩展的方式分配模型。模型参数将被分配并立即在数据并行组中进行分区。如果 remote_device 是 "cpu" 或 "nvme"，模型也将被分配到 CPU/NVMe 内存而不是 GPU 内存。有关更多详细信息，请参阅完整的 ZeRO-3 Init 文档。

     with deepspeed.zero.Init(data_parallel_group=mpu.get_data_parallel_group(),
                          remote_device=get_args().remote_device,
                          enabled=get_args().zero_stage==3):
     model = GPT2Model(num_tokentypes=0, parallel_output=True)

  2. 收集嵌入权重以进行初始化。DeepSpeed 将在模块的构造函数中以及其正向和反向传播过程中自动收集模块的参数。但是，额外的访问必须与 DeepSpeed 协调，以确保参数数据被收集并随后分区。如果张量被修改，modifier_rank 参数也应被使用，以确保所有秩对数据具有一致的视图。有关更多详细信息，请参阅完整的 GatheredParameters 文档。

     self.position_embeddings = torch.nn.Embedding(...)
 with deepspeed.zero.GatheredParameters(self.position_embeddings.weight,
                                        modifier_rank=0):
     # Initialize the position embeddings.
     self.init_method(self.position_embeddings.weight)

 ...

 self.tokentype_embeddings = torch.nn.Embedding(...)
 with deepspeed.zero.GatheredParameters(self.tokentype_embeddings.weight,
                                     modifier_rank=0):
     # Initialize the token-type embeddings.
     self.init_method(self.tokentype_embeddings.weight)

以内存为中心的分块

ZeRO-Infinity 包含一个用于 Linear 层的替代方案，它进一步减少了内存。我们可选择地对每个 Transformer 层中找到的模型并行线性层进行分块。请注意，模型并行和分块可以通过在构建层时指定相应的基类来组合。 deepspeed.zero.TiledLinear 模块利用 ZeRO-3 的数据获取和释放模式，通过将大型操作符分解为可以顺序执行的较小块来减少工作内存需求。

我们包含了 Megatron-LM 的 ParallelMLP 中的一个示例更改。transformer.py 中还有三个模型并行层也类似处理。

Megatron-LM 的模型并行层具有一种特殊形式，其中层的加性 bias 被延迟，而是从 forward() 返回，以便与后续操作符融合。DeepSpeed 的 deepspeed.zero.TiledLinearReturnBias 类是 TiledLinear 的子类，它简单地转发返回的 bias 参数而不进行累加。

@@ -1,6 +1,9 @@
-self.dense_h_to_4h = mpu.ColumnParallelLinear(
+self.dense_h_to_4h = deepspeed.zero.TiledLinearReturnBias(
     args.hidden_size,
     4 * args.hidden_size,
+    in_splits=args.tile_factor,
+    out_splits=4*args.tile_factor,
+    linear_cls=mpu.ColumnParallelLinear,
     gather_output=False,
     init_method=init_method,
     skip_bias_add=True)

请注意，我们根据 input_size 和 output_size 成比例地缩放 in_splits 和 out_splits。这会产生固定大小的块 [hidden/tile_factor, hidden/tile_factor]。

注册外部参数

已弃用： DeepSpeed 版本 0.3.15 引入了自动外部参数注册，此步骤不再需要。

提取权重

如果您需要将预训练权重从 Deepspeed 中取出，以下是获取 fp16 权重的方法：

  • 在 ZeRO-2 下，state_dict 包含 fp16 模型权重，这些权重可以正常使用 torch.save 保存。
  • 在 ZeRO-3 下，state_dict 仅包含占位符，因为模型权重已在多个 GPU 上分区。如果您想获取这些权重，请启用：
    "zero_optimization": {
        "stage3_gather_16bit_weights_on_model_save": true
    },

然后使用以下方式保存模型：

            if self.deepspeed:
                self.deepspeed.save_16bit_model(output_dir, output_file)

由于它需要将权重整合到单个 GPU 上，因此可能会很慢且占用大量内存，所以仅在需要时使用此功能。

请注意，如果 stage3_gather_16bit_weights_on_model_save 为 False，则不会保存任何权重（同样，因为 state_dict 中没有它们）。您也可以使用此方法保存 ZeRO-2 权重。

如果您想获取 fp32 权重，我们提供了一个特殊脚本，可以进行离线整合。它不需要配置文件或 GPU。以下是其用法示例：

$ cd /path/to/checkpoint_dir
$ ./zero_to_fp32.py . pytorch_model.bin
Processing zero checkpoint at global_step1
Detected checkpoint of type zero stage 3, world_size: 2
Saving fp32 state dict to pytorch_model.bin (total_numel=60506624)

当您保存检查点时，zero_to_fp32.py 脚本会自动创建。

注意：目前此脚本使用的内存（通用 RAM）是最终检查点大小的 2 倍。

或者，如果您有足够的空余 CPU 内存，并且希望模型更新为其 fp32 权重而不是获取文件，您可以在训练结束时执行以下操作：

    from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint
    fp32_model = load_state_dict_from_zero_checkpoint(deepspeed.module, checkpoint_dir)

请注意，该模型将适合保存，但不再适合继续训练，并且需要重新执行 deepspeed.initialize()。

如果您只想要 state_dict，您可以执行：

    from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
    state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir)

恭喜！您已完成 ZeRO 教程。

更新日期：