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
  }
}

请注意已添加以支持 1-bit LAMB 功能的新参数 freeze_stepcuda_awarecomm_backend_namecoeff_betafactor_maxfactor_minfactor_threshold

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_maxfactor_minfactor_threshold 用于在压缩阶段规范冻结基缩放系数的自适应缩放。factor_maxfactor_min 是缩放系数上限/下限。factor_threshold 定义了缩放系数在步骤之间可以波动多少的阈值。

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

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

注意! 1-bit LAMB 依赖于压缩误差补偿机制,以在压缩阶段保持收敛速度。加载检查点时,我们实际上重置了压缩误差,原因有 3 点: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/下提供示例脚本。有 3 组脚本对应于基于 NCCL 的实现、基于以太网系统的基于 MPI 的实现以及基于 InfiniBand 系统的基于 MPI 的实现。对于基于 MPI 的实现,我们在使用 deepspeed 或 mpirun 启动时提供两个示例脚本。

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

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

2.3 BERT 预训练的性能结果

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

更新: