通过 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 比特 Adam 优化器相比,0/1 Adam 提供了一种更灵活的方式,通过自适应方差状态冻结来使用压缩通信。此外,它允许计算节点使用一种称为 1 比特同步的技术,在训练期间跳过通信轮次,而不会影响收敛速度。我们有一篇论文提供了技术细节,包括算法、系统实现和评估。
为了说明 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),NCCL 2.8.3 未被 PyTorch 正式支持。我们使用的解决方案是通过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
启动器的 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
,以便可以加载它来执行 float 和 Byte 数据类型之间的高性能打包和解包,这在单比特算法中被利用。示例可以在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
}
}
请注意已添加以支持 0/1 Adam 功能的新参数var_freeze_step
、var_update_scaler
、local_step_scaler
、local_step_clipper
、cuda_aware
和comm_backend_name
var_freeze_step 是更新方差的最新步骤。使用0/1 Adam 论文中的符号,它表示 $\max{i |
i \in \mathcal{T}_v}$。请注意,这与 1 比特 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 感知通信。此功能仅在具有 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
。
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 位 Adam,我们还需要重置本地步长缓冲区。因为如果检查点由不同数量的节点(GPU)加载,本地步长缓冲区可能会无法捕获训练动态。
2. 使用 0/1 Adam 进行 BERT 预训练
有关数据下载和预处理,请参阅 BERT 预训练教程。
2.1 使用 DeepSpeed 和 0/1 Adam 运行预训练
我们在 DeepSpeedExamples/bing_bert/01_adam/ 下提供示例脚本。有 3 组脚本,分别对应基于 NCCL 的实现、基于 MPI 在以太网系统上的实现,以及基于 MPI 在 InfiniBand 系统上的实现。对于基于 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 微调
我们还提供了针对 BERT 预训练检查点的微调脚本,用于在 GLUE 任务 上进行微调。脚本位于 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 中。