1比特LAMB：通信高效的大规模大批量训练，同时保持LAMB的收敛速度

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

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

为了说明 1比特LAMB 优化器的优势和用法，我们以 BERT 预训练任务为例。有关此任务的更多详细信息，请参阅教程。

1. 概述

1.1 安装 DeepSpeed 的先决条件

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

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

1.2 1比特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_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比特LAMB 的示例如下

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

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

1.3 1比特LAMB 算法

1比特LAMB 算法的详细描述可以在我们的论文中找到。

1.4 1比特LAMB 配置

1比特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
  }
}

请注意，为了支持 1比特LAMB 功能，新增了参数 freeze_step、cuda_aware、comm_backend_name、coeff_beta、factor_max、factor_min 和 factor_threshold。

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

cuda_aware 用于基于 MPI 的实现，以指示底层 MPI 库支持 CUDA-Aware 通信。此功能仅在具有 InfiniBand 互连和支持 CUDA-Aware 的 MPI 库（如 MVAPICH2-GDR 或使用 CUDA-Aware 支持构建的 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比特压缩无法表示精确的零，如果参数在训练期间具有恒定的零梯度，则压缩误差将持续在动量中累积。例如，对于 BERT 预训练序列长度 128，bert.embeddings.position_embeddings.weight 在其梯度和动量中，第 129 到 512 行具有恒定零值，因为它仅学习到序列长度 128，而模型支持最大序列长度 512。因此，在 1比特LAMB 中，我们添加了对动量掩码的支持，供用户指定那些梯度中具有恒定精确零值的参数。有关如何配置此动量掩码，请参阅示例脚本。需要注意的是，我们不使用检查点中保存的动量掩码，因为此掩码在训练期间可能会改变（例如，BERT 序列长度 128 和 512 需要不同的掩码）。因此，您必须在训练脚本中每次提供此掩码。

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

2. 使用 1比特LAMB 进行 BERT 预训练

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

2.1 使用 DeepSpeed 和 1比特LAMB 运行预训练

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

2.2 启用 DeepSpeed 和 1比特LAMB 进行 BERT 预训练的配置

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

2.3 BERT 预训练的性能结果

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

更新时间：