1 位 Adam：通信量减少 5 倍，训练速度提高 3.4 倍

注意： 2022 年 3 月 7 日，我们发布了 0/1 Adam，它是一个新的通信效率高的 Adam 优化器，部分遵循 1 位 Adam 的设计。与下面描述的 1 位 Adam 相比，0/1 Adam 提供了更好的通信效率，并在 BERT、GPT-2 和 ImageNet 等不同任务上实现了相同的最终模型质量。因此，我们建议您首先尝试 0/1 Adam (教程)，然后在 0/1 Adam 无法在您的任务中提供基线 Adam 收敛的情况下尝试 1 位 Adam。

注意： 本教程已于 2021 年 3 月 4 日更新，以反映 1 位 Adam v2。更改包括：1) 基于 NCCL 的实现，与基于 MPI 的实现相比，它提供了更好的性能和可用性。2) 添加对动量掩码的支持，用于在训练期间梯度恒为零的参数。3) 错误修复。请参见下面的详细信息。

注意！ 1) 基于 NCCL 的实现需要 PyTorch >= 1.8（如果您有 64 个或更多 GPU，则需要 NCCL >= 2.8.3）。请参见下面的详细信息。2) 尽管 1 位 Adam 与 FP16 和 FP32 兼容，但目前我们只验证了混合精度/FP16 训练下的收敛性。3) 目前基于 MPI 的实现与管道并行不兼容。4) 频繁加载检查点可能会损害 1 位 Adam 的收敛性。请参见下面的详细信息。

在本教程中，我们将介绍 DeepSpeed 中的 1 位 Adam 优化器。1 位 Adam 可以提高通信受限集群上的模型训练速度，特别是对于通信密集型的大型模型，它可以将整体通信量减少多达 5 倍。有关 1 位 Adam 算法的详细描述、其在 DeepSpeed 中的实现以及性能评估，请参阅我们的 博客文章。我们还有一篇 论文，其中提供了最完整的详细信息，包括算法、系统实现、理论分析以及更多评估。

为了说明 1 位 Adam 优化器在 DeepSpeed 中的优势和用法，我们以以下两个训练任务为例

1. BingBertSQuAD 微调
2. BERT 预训练

有关这些任务的更多详细信息，请参阅有关 BingBertSQuAD 微调 和 BERT 预训练 的教程文章。

1. 概述

1.1 安装 DeepSpeed 的先决条件

如果您还没有 DeepSpeed 存储库的副本，请立即克隆它并签出包含 BingBertSQuAD 和 BERT 预训练示例的 DeepSpeedExamples 子模块。

git clone https://github.com/microsoft/DeepSpeed
cd DeepSpeed
git submodule update --init --recursive
cd DeepSpeedExamples/

1.2 1 位 Adam 的先决条件

1.2.1 (v2 中的新功能) 基于 NCCL 的实现

在 1 位 Adam v2 中，我们引入了使用 PyTorch 分布式 NCCL 后端进行压缩通信的新系统实现。这极大地提高了可用性，因为 NCCL 与 PyTorch 分布式集成在一起。我们新的基于 NCCL 的实现的性能也优于我们早期的基于 MPI 的实现（对于基于以太网的系统）以及基于 InfiniBand 的系统。因此，我们强烈建议用户选择此实现。

注意！ 此基于 NCCL 的实现需要 PyTorch >= 1.8。如果您有 64 个或更多 GPU，则还需要 NCCL >= 2.8.3，以避免某些 NCCL 运行时错误。目前（2021/03/16），NCCL 2.8.3 未被 PyTorch 正式支持。我们使用的解决方案是通过 LD_PRELOAD 侵入 NCCL 2.8.3：1) 安装 NCCL 2.8.3。这在 CUDA 11 系统上对我们有效：apt-get install -y libnccl2=2.8.3-1+cuda11.0 libnccl-dev=2.8.3-1+cuda11.0。2) 将 LD_PRELOAD 设置为库路径。这在我们的系统上有效：LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libnccl.so.2.8.3。要确认 LD_PRELOAD 是否有效，您可以在 NCCL 日志中查看它使用的版本（如果您有 NCCL_DEBUG=INFO），它应该显示：NCCL 版本 2.8.3+cuda11.0。

1.2.2 基于 MPI 的实现

对于此实现，我们依赖于消息传递接口 (MPI) 来实现高级通信原语。

我们在 DeepSpeed Docker 镜像中打包了必要的依赖项。但是，如果您使用的是其他构建系统，请在您的系统上安装 MPI 和 mpi4py。要安装先决条件，请运行

pip install deepspeed[1bit_adam]

我们已使用 MVAPICH2-GDR 库测试了 CUDA 感知 MPI 通信。但是，任何 CUDA 感知通信库（包括 OpenMPI）都应该可以与这些示例正常配合使用。

使用 deepspeed 启动器启动 1 位 Adam 的示例命令如下

deepspeed --launcher=[mvapich|openmpi] script.py

请注意，对于 1 位 Adam 的基于 MPI 的实现，使用 deepspeed 启动器时，需要 --launcher=[mvapich|openmpi] 标志。

或者，也可以使用标准 mpirun 启动器，如下所示

mpirun -np [#processes] -ppn [#GPUs on each node] -hostfile [hostfile] [MPI flags] python [training_script.py]

1.2.3 压缩实现

此后端提供了一种方法来抽象一位优化器的通用部分，并使用 DeepSpeed 自定义操作构建器实现加速器相关的部分。要使用此 CompressedBackend，您应该确保您当前的加速器支持 PackbitsBuilder，以便它可以加载以在浮点和字节数据类型之间执行高性能打包和解包，这在一位算法中被使用。可以在 Deepspeed/op_builder/xpu/packbits.py 中找到示例。

此方法不需要基于 NCCL 或 MPI 的通信库。它会自动使用您在 deepspeed/comm 中选择的加速器默认通信库。

1.3 1 位算法

有关 1 位算法的详细描述，请参阅我们的 博客文章 和我们的 论文。

1.4 1 位 Adam 的配置

可以通过如下设置优化器配置选项来使用 1 位 Adam 功能。下面显示了一个示例 json 配置文件。

{
  "train_batch_size": 4096,
  "train_micro_batch_size_per_gpu": 16,
  "optimizer": {
    "type": "OneBitAdam",
    "params": {
      "lr": 4e-4,
      "freeze_step": 23000,
      "cuda_aware": false,
      "comm_backend_name": "nccl"
    }
  },
  "fp16": {
    "enabled": true,
  }
}

请注意，已添加了三个新参数 freeze_step、cuda_aware 和 comm_backend_name 来支持 1 位 Adam 功能。

freeze_step 是在将 1 位压缩应用于通信之前预热步骤的数量。为了确定预热步骤的数量，一种策略是将给定模型的总训练步骤的 15% 到 25% 设置为预热步骤（这与 Adam 的方差/二阶矩项有关。有关详细分析，请参阅我们的 论文）。如果它产生了预期结果，则可以尝试通过系统地减少步骤来获得更多性能。将来，我们计划引入一个阈值，该阈值可以自动搜索和决定不同模型的预热步骤数量。下面的示例已针对预热步骤数量进行了调整。freeze_step 参数已在相应的运行脚本中设置为我们找到的最佳数字。

cuda_aware 用于基于 MPI 的实现，以指示底层 MPI 库是否支持 CUDA 感知通信。此功能仅在具有 InfiniBand 互连和 CUDA 感知 MPI 库（如 MVAPICH2-GDR 或使用 CUDA 感知支持构建的 OpenMPI）的系统上受支持。将 cuda_aware 设置为 False 可以在基于以太网的系统上进行训练。但是，通信将在通信之前和之后使用发送方和接收方内存复制在 CPU 和 GPU 缓冲区之间进行。

(v2 中的新功能) comm_backend_name 用于指示要使用哪个后端实现。您可以通过将 comm_backend_name 设置为“nccl”、“mpi”或“compressed”来在基于 NCCL、基于 MPI 和压缩的实现之间进行选择。使用基于 NCCL 的实现时，无需设置 cuda_aware。

1.4.1 (v2 中的新功能) 针对梯度恒为零的参数的动量掩码

由于 1 位压缩无法表示精确的零，因此如果参数在训练期间梯度恒为零，则压缩误差将在动量中不断累积。例如，对于 BERT 预训练 seq 长度 128，bert.embeddings.position_embeddings.weight 在其梯度和动量中对于第 129 行到第 512 行有恒定的零，因为它只学习到 seq 长度 128，而模型支持的 seq 长度为 512。因此，在 1 位 Adam v2 中，我们添加了对动量掩码的支持，供用户指定那些在梯度中具有恒定精确零的参数。有关如何配置此动量掩码，请参阅 示例脚本。需要注意的是，我们不使用检查点中保存的动量掩码，因为此掩码可能会在训练期间发生变化（例如，BERT seqlen 128 和 512 需要不同的掩码）。因此，您必须每次都在训练脚本中提供此掩码。

注意！ 1 位 Adam 依赖于压缩误差补偿机制来维护压缩阶段的收敛速度。加载检查点时，我们实际上会重置压缩误差，原因有 3 个：1) 每个 GPU 上的 worker 和 server 误差是不同的，因此在当前实现中，只有 rank 0 的误差保存在检查点中。因此，我们必须重置误差。如果我们想正确保存它们，我们需要 O(num_gpu*model_size) 的内存才能收集所有误差，这是一个非常大的内存需求。可以以分布式方式保存它们，但这将使检查点保存/加载变得更加复杂。2) 即使我们能够正确保存压缩误差，您也需要具有相同数量的 GPU 才能正确加载它们。3) 我们在 BERT 预训练上验证了，偶尔在加载检查点时重置压缩误差不会影响收敛性。但是，请避免频繁加载检查点，这可能会破坏误差补偿机制，从而影响收敛性。

2. 使用 1 位 Adam 进行 BingBertSQuAD 微调

• 下载 SQuAD 数据集
  • 训练集：train-v1.1.json
  • 验证集：dev-v1.1.json
• 下载 HuggingFace 检查点和配置文件
  • bert-large-uncased-whole-word-masking
  • bert json 配置

您也可以使用来自 DeepSpeed、HuggingFace 或 TensorFlow 的预训练 BERT 模型检查点来运行微调。

注意：有关加载检查点、参数解析、初始化、前向传播、反向传播、权重更新和评估的详细信息，请参阅BingBertSQuAD 微调教程。

2.1 使用 DeepSpeed 和 1 位 Adam 运行 BingBertSQuAD

我们在DeepSpeedExamples/BingBertSquad/1-bit_adam/下提供示例脚本。有 3 组脚本对应于基于 NCCL 的实现、以太网系统上的基于 MPI 的实现以及 InfiniBand 系统上的基于 MPI 的实现。对于基于 MPI 的实现，我们在使用 deepspeed 或 mpirun 启动时都提供了示例脚本。

2.2 使用 DeepSpeed 和 1 位 Adam 启用的 BingBertSQuAD 配置

The deepspeed_onebitadam_bsz96_config.json 文件允许用户指定 DeepSpeed 选项，包括批大小、微批大小、优化器、学习率和其他参数。在运行 nvidia_run_squad_deepspeed.py 时，除了使用 --deepspeed 标志启用 DeepSpeed 之外，还必须使用 --deepspeed_config deepspeed_onebitadam_bsz96_config.json 指定相应的 DeepSpeed 配置文件。

表 1 显示了我们在实验中使用的微调配置。

参数	值
总批大小	96
每个 GPU 的训练微批大小	3
优化器	“OnebitAdam”
学习率	3e-5
序列长度	384
权重衰减	0.0
历元计数	2
freeze_step	400
comm_backend_name	“nccl”

表 1. 微调配置

2.3 BingBertSQuAD 微调的性能结果

准确率：结果总结在下面的表格中。总批大小设置为 96，在 32 个 GPU 上进行 2 个历元的训练。尝试了一组参数（种子和学习率），并选择了最佳参数。我们将学习率固定为 3e-5。下表显示了我们取得的 F1 和 EM 分数，它们与 HuggingFace 结果 相当或更好。

案例	模型	精度	EM	F1
HuggingFace	Bert-large-uncased-whole-word-masking	FP16	87.26	93.32

训练速度和可扩展性

SQuAD 微调的性能结果可以从我们的博客文章和我们的论文中看到。

3. 使用 1 位 Adam 进行 BERT 预训练

有关数据下载和预处理，请参阅BERT 预训练教程。

3.1 使用 DeepSpeed 和 1 位 Adam 运行预训练

我们在DeepSpeedExamples/bing_bert/1-bit_adam/下提供示例脚本。有 3 组脚本对应于基于 NCCL 的实现、以太网系统上的基于 MPI 的实现以及 InfiniBand 系统上的基于 MPI 的实现。对于基于 MPI 的实现，我们在使用 deepspeed 或 mpirun 启动时都提供了示例脚本。

3.2 使用 DeepSpeed 和 1 位 Adam 启用的 BERT 预训练配置

The deepspeed_bsz4k_onebit_config_seq128_*.json 文件允许用户指定 DeepSpeed 选项，包括批大小、微批大小、优化器、学习率和其他参数。

以下是使用 1 位 Adam 优化器以 128 的序列长度运行 BERT-large 预训练的 DeepSpeed 配置文件。

{
  "train_batch_size": 4096,
  "train_micro_batch_size_per_gpu": 16,
  "steps_per_print": 100,
  "prescale_gradients": false,
  "optimizer": {
    "type": "OneBitAdam",
    "params": {
      "lr": 4e-4,
      "weight_decay": 0.01,
      "bias_correction": false,
      "freeze_step": 23000,
      "comm_backend_name": "nccl"
    }
  },
  "gradient_clipping": 1.0,
  "fp16": {
    "enabled": true,
    "loss_scale": 0,
    "initial_scale_power": 16
  }
}

上述文件适用于 BERT-large。对于 BERT-base 训练（序列长度 128），建议的 freeze_step 为 16000。对于序列 512 预训练，我们建议对 BERT-base 和 BERT-large 都使用 freeze_step 为 1500。并确保按上述说明正确设置 comm_backend_name 和 cuda_aware。

3.3 BERT 预训练的性能结果

BERT 预训练的性能结果可以从我们的博客文章和我们的论文中看到。