通过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-bit Adam优化器相比，0/1 Adam通过自适应方差状态冻结提供了一种更灵活的使用压缩通信的方式。此外，它允许计算节点在训练期间使用一种称为1-bit sync的技术跳过通信轮次，而不会影响收敛速度。我们有一篇论文，提供了包括算法、系统实现和评估在内的技术细节。

为了说明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_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启动器的0/1 Adam的示例启动命令如下：

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

请注意，对于基于MPI的0/1 Adam实现，在使用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 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-bit 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-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。

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-bit Adam一样重置压缩误差外，我们还需要重置本地步长缓冲区。因为如果检查点是由不同数量的节点（GPU）加载的，本地步长缓冲区可能无法捕捉到训练动态。

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

关于数据下载和预处理的详细信息，请参考BERT预训练教程。

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

我们在DeepSpeedExamples/bing_bert/01_adam/下提供了示例脚本。有三组脚本分别对应于基于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中。