零冗余优化器

如果您还没有这样做，我们建议您在继续本教程之前阅读关于入门和Megatron-LM GPT-2的DeepSpeed教程。

在本教程中，我们将应用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 包含了 infinity offload engine 以形成 ZeRO-Infinity (paper)，它可以卸载到 CPU 和 NVMe 内存，从而实现巨大的内存节省。

训练环境

我们使用DeepSpeed Megatron-LM GPT-2代码进行此练习。你可以逐步学习Megatron-LM 教程以熟悉代码。我们将在本教程中使用NVIDIA Tesla V100-SXM3 Tensor Core GPUs（32GB内存）来训练模型。

启用ZeRO优化

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

训练一个15亿参数的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内存的很大一部分。通过使用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 第二阶段优化进一步增加了可以使用数据并行训练的模型大小。我们通过在32个V100 GPU上训练一个具有100亿参数的模型来展示这一点。

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

       --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以减少反向传播期间的内存碎片。这些优化选项的完整描述可在此处here找到。通过这些更改，我们现在可以启动训练运行。

以下是训练日志的截图：

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

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

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

使用ZeRO-Infinity将任务卸载到CPU和NVMe

ZeRO-Infinity 使用 DeepSpeed 的 infinity 卸载引擎将完整的模型状态卸载到 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)

请注意，我们按比例缩放in_splits和out_splits与input_size和output_size。这导致了固定大小的瓦片[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教程。

更新: