1-bit Adam：通信量减少高达 5 倍，训练速度提升高达 3.4 倍

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

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

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

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

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

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-bit Adam 的先决条件

1.2.1（v2 中新增）基于 NCCL 的实现

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

注意！此基于 NCCL 的实现需要 PyTorch >= 1.8。当您有 64 个或更多 GPU 时，它还需要 NCCL >= 2.8.3 以避免某些 NCCL 运行时错误。目前（2021/03/16）PyTorch 不正式支持 NCCL 2.8.3。我们使用的解决方案是通过 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-bit Adam 的示例命令如下所示

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

请注意，对于 1-bit 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-bit 算法

1-bit 算法的详细描述可以在我们的博客文章和我们的论文中看到。

1.4 1-bit Adam 的配置

可以通过如下设置优化器配置选项来使用 1-bit 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-bit Adam 功能。

freeze_step 是在将 1-bit 压缩应用于通信之前预热步骤的数量。为了确定预热步骤的数量，一种策略是将给定模型的总训练步骤的 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-bit 压缩无法表示精确的零，因此如果参数在训练期间具有恒定零梯度，则压缩误差会不断累积在动量中。例如，对于 BERT 预训练序列长度 128，bert.embeddings.position_embeddings.weight 在其第 129 行到第 512 行的梯度和动量中具有恒定零，因为它仅学习到序列长度 128，而模型支持的序列长度高达 512。因此，在 1-bit Adam v2 中，我们添加了动量掩码的支持，供用户指定其梯度中具有恒定精确零的参数。请参阅示例脚本，了解如何配置此动量掩码。需要注意的是，我们不使用检查点中保存的动量掩码，因为此掩码可能会在训练期间发生更改（例如，BERT 序列长度 128 和 512 需要不同的掩码）。因此，您必须在每次训练脚本中提供此掩码。

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

2. 使用 1-bit 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-bit Adam 运行 BingBertSQuAD

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

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

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-bit Adam 进行 BERT 预训练

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

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

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

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

deepspeed_bsz4k_onebit_config_seq128_*.json 文件使用户能够根据批次大小、微批次大小、优化器、学习率和其他参数指定 DeepSpeed 选项。

以下是使用 1-bit 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 预训练的性能结果。