1-bit LAMB：使用LAMB收敛速度进行通信高效的大规模大批量训练

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

在本教程中，我们介绍了 DeepSpeed 的 1-bit LAMB 优化器，它可以使用 LAMB 的收敛速度进行通信高效的大规模大批量训练。1-bit LAMB 可以通过将整体通信量减少高达 4.6 倍，从而提高通信受限集群上的模型训练速度，特别是对于通信密集型大型模型。我们还有一篇论文，其中提供了技术细节，包括算法、系统实现和评估。

为了说明 1-bit LAMB 优化器的优势和用法，我们以 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 1-bit LAMB 的先决条件

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 启动器启动 1-bit LAMB 的示例命令如下所示

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

请注意，对于基于 MPI 的 1-bit LAMB 实现，在使用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，以便可以加载它来执行浮点数和字节数据类型之间的高性能打包和解包，这在单比特算法中使用。示例可以在Deepspeed/op_builder/xpu/packbits.py 中找到。此方法不需要基于 NCCL 或 MPI 的通信库。它将自动使用deepspeed/comm 中加速器选择的默认通信库。

1.3 1-bit LAMB 算法

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

1.4 1-bit LAMB 的配置

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

{
  "train_batch_size": 65536,
  "train_micro_batch_size_per_gpu": 64,
  "optimizer": {
    "type": "OneBitLamb",
    "params": {
      "lr": 11e-3,
      "max_coeff": 0.3,
      "min_coeff": 0.01,
      "freeze_step": 1000,
      "cuda_aware": false,
      "comm_backend_name": "nccl",
      "coeff_beta": 0.9,
      "factor_max": 4.0,
      "factor_min": 0.5,
      "factor_threshold": 0.1
    }
  },
  "gradient_clipping": 1.0,
  "fp16": {
    "enabled": true,
    "loss_scale": 0,
    "initial_scale_power": 16
  }
}

请注意已添加的新参数freeze_step、cuda_aware、comm_backend_name、coeff_beta、factor_max、factor_min 和factor_threshold，以支持 1-bit LAMB 功能

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

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。

coeff_beta 用于在预热阶段计算 LAMB 缩放系数的移动平均值。然后，此移动平均值在压缩阶段用作冻结的基本缩放系数。

factor_max、factor_min 和factor_threshold 用于在压缩阶段规范冻结的基本缩放系数的自适应缩放。factor_max 和factor_min 是缩放因子的上限/下限。factor_threshold 定义缩放因子在步骤之间可以波动多少的阈值。

1.4.1 梯度始终为零的参数的动量掩码

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

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

2. 使用 1-bit LAMB 进行 BERT 预训练

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

2.1 使用 DeepSpeed 和 1-bit LAMB 运行预训练

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

2.2 使用 DeepSpeed 和启用 1-bit LAMB 的 BERT 预训练配置

deepspeed_bsz64k_onebitlamb_config_seq128_*.json 和 deepspeed_bsz32k_onebitlamb_config_seq512_*.json 文件允许用户根据批大小、微批大小、优化器、学习率和其他参数指定 DeepSpeed 选项。在这些文件中，我们包含了用于复现我们论文中实验的已调优超参数。

2.3 BERT 预训练的性能结果

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