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_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 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位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位压缩应用于通信之前的热身步骤数。为了确定热身步骤的数量，一种策略是为给定模型设置总训练步骤的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-bit LAMB 依赖于压缩误差补偿机制来保持压缩阶段的收敛速度。在加载检查点时，我们实际上会重置压缩误差，原因有三：1) 每个 GPU 上的工作器和服务器误差是不同的，因此在当前实现中，只有 rank 0 的误差被保存在检查点中。因此我们必须重置误差。如果我们想正确保存它们，我们需要 O(num_gpu*model_size) 的内存来收集所有误差，这是一个非常大的内存需求。可以以分布式方式保存它们，但这会使检查点的保存/加载变得更加复杂。2) 即使我们能够正确保存压缩误差，你也需要有完全相同数量的 GPU 才能正确加载它们。3) 我们在 BERT 预训练中验证了在检查点加载时偶尔重置压缩误差不会影响收敛。然而，请避免频繁加载检查点，这可能会破坏误差补偿机制，从而影响收敛。

2. 使用1位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预训练的性能结果

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

更新: