通过 0/1 Adam 优化器最大化大规模训练的通信效率

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

在本教程中，我们介绍了 DeepSpeed 的 0/1 Adam 优化器，它可以提高在通信受限集群上模型训练的速度，特别是对于通信密集型的大型模型。例如，它能够将 BERT-large 预训练的整体通信量减少高达 26 倍，而不会影响端到端模型的准确性。与 1 比特 Adam 优化器相比，0/1 Adam 提供了一种更灵活的方式来通过自适应方差状态冻结使用压缩通信。此外，它允许计算节点使用称为 1 比特同步的技术跳过训练期间的通信轮次，而不会影响收敛速度。我们有一篇论文，其中提供了技术细节，包括算法、系统实现和评估。

为了说明 0/1 Adam 优化器的优势和用法，我们使用 BERT 预训练任务作为示例。有关此任务的更多详细信息，请参阅教程。

1. 概述

1.1 安装 DeepSpeed 的先决条件

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

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

1.2 0/1 Adam 的先决条件

1.2.1 基于 NCCL 的实现

在 DeepSpeed 中，我们引入了使用 PyTorch 分布式 NCCL 后端的压缩通信的系统实现。此实现提供了比下面基于 MPI 的实现更好的性能和可用性。因此，我们强烈建议用户选择此实现。

注意！此基于 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 启动器启动 0/1 Adam 的示例命令如下所示

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

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

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

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

1.2.3 压缩实现

此后端提供了一种方法来抽象出一比特优化器的通用部分，并使用 DeepSpeed 自定义操作构建器实现加速器相关的部分。要使用此 CompressedBackend，您应该确保您当前的加速器支持 PackbitsBuilder，以便可以加载它来在浮点数和字节数据类型之间进行高性能打包和解包，这在 1 比特算法中使用。一个示例可以在 Deepspeed/op_builder/xpu/packbits.py 中找到。此方法不需要基于 NCCL 或 MPI 的通信库。它将自动使用您的加速器在 deepspeed/comm 中选择的默认通信库。

1.3 0/1 Adam 算法

0/1 Adam 算法的详细描述可以在我们的论文中看到。

1.4 0/1 Adam 的配置

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

{
  "train_batch_size": 4096,
  "train_micro_batch_size_per_gpu": 16,
  "optimizer": {
    "type": "ZeroOneAdam",
    "params": {
      "lr": 1e-3,
      "weight_decay": 0.01,
      "bias_correction": false,
      "var_freeze_step": 1000,
      "var_update_scaler": 16,
      "local_step_scaler": 1000,
      "local_step_clipper": 16,
      "cuda_aware": false,
      "comm_backend_name": "nccl"
    }
  },
  "gradient_clipping": 1.0,
  "fp16": {
    "enabled": true,
    "loss_scale": 0,
    "initial_scale_power": 16
  }
}

请注意已添加的新参数 var_freeze_step、var_update_scaler、local_step_scaler、local_step_clipper、cuda_aware 和 comm_backend_name 以支持 0/1 Adam 功能

var_freeze_step 是更新方差的最新步骤。使用 0/1 Adam 论文中的符号，它表示 $\max{i i \in \mathcal{T}_v}$。请注意，这与 1 比特 Adam 中的 freeze_step 不同。 var_freeze_step 通常是学习率预热期的最后一步，因此不需要调整。请注意，此超参数是可选的。在实践中，我们可以通过将其设置为足够大的数字（大于步骤总数）来避免调整此参数。在此之后，0/1 Adam 仍然可以享受非平凡的通信减少，而不会影响收敛速度。

var_update_scaler 是更新方差的间隔。请注意，方差的更新策略遵循指数规则。形式上，如果我们将 $k_j$ 表示为第 $j$ 次方差更新发生的步骤，则它遵循 $k_{j+1} - k_j = 2\cdot\exp{\lfloor j/\kappa\rfloor}$（有关详细说明，请参阅0/1 Adam 论文），并且 var_update_scaler 表示该表达式中的 $\kappa$ 因子。在实践中，我们发现它的默认值 (16) 能够很好地处理大多数任务，包括 BERT-Base/Large 预训练、GPT 预训练和 ImageNet 训练。

local_step_scaler 和 local_step_clipper 是 0/1 Adam 中基于学习率的本地步骤策略的两个超参数。形式上，如果我们将 $k_j$ 表示为第 $j$ 次同步在所有工作器之间发生的步骤，则它遵循 $k_{j+1} - k_j = 2\cdot\exp{\min(\lfloor j/\alpha\rfloor, \beta )}$（有关详细说明，请参阅0/1 Adam 论文）。根据此类符号，local_step_scaler 和 local_step_clipper 分别表示 $\alpha$ 和 $\beta$。非正式地，local_step_scaler 决定同步的频率，而 local_step_clipper 表示 0/1 Adam 可以使用的最大本地步骤间隔。学习率策略是 0/1 Adam 中使用的默认策略，并且 local_step_scaler 的值可以预先计算（参见0/1 Adam 论文第 6 节）。我们还可以通过设置这两个超参数来简单地构建其他策略，例如通过设置 local_step_scaler=1 和 local_step_clipper=constant 来设置恒定的本地步骤间隔策略。

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

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

1.4.1 具有恒定零梯度的参数的动量掩码

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

注意！ 0/1 Adam依赖于压缩误差补偿机制来维持压缩阶段的收敛速度。在加载检查点时，除了像1比特Adam一样重置压缩误差外，我们还需要重置本地步骤缓冲区。因为如果检查点由不同数量的节点（GPU）加载，本地步骤缓冲区可能会无法捕获训练动态。

2. 使用 0/1 Adam 进行 BERT 预训练

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

2.1 使用 DeepSpeed 和 0/1 Adam 运行预训练

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

2.2 使用 DeepSpeed 和启用 0/1 Adam 的 BERT 预训练的配置

deepspeed_bsz4k_01adam_config_seq128_*.json和deepspeed_bsz4k_01adam_config_seq512_*.json文件使用户能够根据批大小、微批大小、优化器、学习率和其他参数指定DeepSpeed选项。在这些文件中，我们包含了用于重现我们论文中实验的调整后的超参数。

2.3 BERT 预训练的性能结果

性能结果可以在我们的论文中查看。

2.4 GLUE 微调

我们还提供了针对GLUE任务的BERT预训练检查点的微调脚本。这些脚本位于DeepSpeedExamples/BingBertGlue。 glue_bert_base.json和glue_bert_large.json文件分别使用户能够指定BERT-base和BERT-large检查点的DeepSpeed选项/参数，例如微批大小。目前，我们使用Adam作为GLUE微调的默认优化器，因为微调任务通常使用较小的批大小（约32）并且不需要大型系统。run_glue_bert_base_finetune.sh和run_glue_bert_large_finetune.sh提供了用于启动微调任务的脚本，我们可以在其中修改任务名称、历元数、模型等变量。请注意，要启动微调，我们必须指定检查点的路径，例如：

bash run_glue_bert_base_finetune.sh <path to checkpoint>

0/1 Adam的具体GLUE分数和超参数包含在我们论文的表1中。