零冗余优化器

如果您还没有这样做，我们建议您在逐步完成本教程之前，阅读 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教程以熟悉代码。我们将在本教程中使用NVIDIA Tesla V100-SXM3 Tensor Core GPU（内存为 32GB）对模型进行训练。

启用 ZeRO 优化

要为 DeepSpeed 模型启用 ZeRO 优化，我们只需将zero_optimization键添加到 DeepSpeed JSON 配置中。zero_optimization键的配置旋钮的完整描述在此处提供。

训练一个 15 亿参数的 GPT-2 模型

我们通过展示它如何在八个 V100 GPU 上启用 15 亿参数 GPT-2 模型的数据并行训练来演示 ZeRO 阶段 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，我们可以通过增加数据并行度来进一步减少每个设备的内存消耗。这些内存节省可以用来增加模型大小和/或批次大小。相反，仅使用数据并行则无法获得此类好处。

训练一个 100 亿参数的 GPT-2 模型

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

首先，我们需要配置一个启用了激活检查点的 100 亿参数模型。这可以通过对 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以减少反向传递期间的内存碎片。这些优化旋钮的完整描述在此处提供。通过这些更改，我们现在可以启动训练运行。

这是训练日志的屏幕截图

这是 nvidia-smi 的屏幕截图，显示了训练期间的 GPU 活动

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

ZeRO-3 是 ZeRO 的第三阶段，它将完整模型状态（即权重、梯度和优化器状态）进行分区，以使内存节省量与数据并行度线性扩展。可以在 JSON 配置中启用 ZeRO-3。这些配置的完整描述在此处提供。

将数据卸载到 CPU 和 NVMe 中，使用 ZeRO-Infinity

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-Offload 中包含了卸载功能，ZeRO-Offload 是一个用于在 ZeRO-2 中将优化器和梯度状态卸载到 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）是最终检查点大小的两倍。

或者，如果您有足够的备用 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 教程。

更新：