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/deepspeedai/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年3月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_DEBUG=INFO，可以在 NCCL 日志中看到它使用的版本，应显示：NCCL version 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比特压缩之前进行热身（warm up）的步数。为了确定热身步数，一种策略是为给定模型设置总训练步数的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 预训练序列长度为128的情况，bert.embeddings.position_embeddings.weight 的第129到512行的梯度和动量始终为零，因为它只学习到序列长度128，而模型支持的序列长度可达512。因此，在 1-bit Adam v2 中，我们添加了动量掩码支持，供用户指定那些梯度始终为零的参数。有关如何配置此动量掩码，请参见示例脚本。需要注意的是，我们不使用检查点中保存的动量掩码，因为此掩码在训练期间可能会改变（例如，BERT 序列长度128和512需要不同的掩码）。因此，您必须在每次训练脚本中提供此掩码。

注意！ 1-bit Adam 依赖于压缩误差补偿机制来维持压缩阶段的收敛速度。加载检查点时，我们实际上会重置压缩误差，原因有三：1) 每个 GPU 上的工作器和服务端误差是不同的，因此在当前实现中，只有 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 实现以及基于 InfiniBand 系统的 MPI 实现。对于基于 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
Epoch 计数	2
freeze_step	400
comm_backend_name	“nccl”

表1. 微调配置

2.3 BingBertSQuAD 微调的性能结果

准确性： 结果总结在下表中。总批大小设置为96，在32个 GPU 上训练2个 Epoch。尝试了一组参数（种子和学习率），并选择了最佳的。我们将学习率固定为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 实现以及基于 InfiniBand 系统的 MPI 实现。对于基于 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 都使用1500的 freeze_step。并且请确保如上所述正确设置 comm_backend_name 和 cuda_aware。

3.3 BERT 预训练的性能结果

BERT 预训练的性能结果可参见我们的博客文章和我们的论文。