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_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-Aware MPI通信。然而，包括OpenMPI在内的任何CUDA-Aware通信库都应该能很好地与这些示例一起工作。

使用deepspeed启动器进行1-bit Adam的示例启动命令如下：

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

请注意，对于基于MPI的1-bit Adam实现，在使用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-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位压缩应用于通信之前的热身步骤数。为了确定热身步骤的数量，一种策略是为给定模型设置总训练步骤的15-25%（这与Adam的方差/第二矩项有关。详见我们的论文中的详细分析）。如果它提供了预期的结果，可以尝试通过系统地减少步骤来提取更多性能。未来，我们计划引入一个阈值，可以自动搜索并决定不同模型的热身步骤数。以下示例已经针对热身步骤数进行了调整。freeze_step 参数已经设置为我们在相应运行脚本中找到的最佳数值。

cuda_aware 用于基于MPI的实现，以指示底层MPI库支持CUDA-Aware通信。此功能仅在具有InfiniBand互连和CUDA-Aware MPI库（如MVAPICH2-GDR或构建时支持CUDA-Aware的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位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 config

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

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

2.1 使用DeepSpeed和1-bit Adam运行BingBertSQuAD

我们在DeepSpeedExamples/training/BingBertSquad/1-bit_adam/下提供了示例脚本。有三组脚本分别对应基于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
训练轮数	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/下提供了示例脚本。有三组脚本分别对应基于NCCL的实现、在以太网系统上基于MPI的实现，以及在InfiniBand系统上基于MPI的实现。对于基于MPI的实现，我们提供了使用deepspeed或mpirun启动时的示例脚本。

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

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

以下是使用1-bit Adam优化器运行BERT-large预训练（序列长度为128）的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预训练的性能结果可以从我们的博客文章和我们的论文中看到。