DeepSpeed 1-bit Adam 训练优化：5 倍通信压缩与 3.4 倍训练加速的完整配置指南

推理引擎大模型【免费下载链接】FlexGenRunning large language models on a single GPU for throughput-oriented scenarios.项目地址https://gitcode.com/gh_mirrors/fl/FlexGen点击查看免费下载1-bit Adam 是 DeepSpeed 中面向通信受限集群的高效训练优化器通过将分布式训练中梯度同步的通信量压缩至原来的约 1/5可在通信密集型大模型训练场景中显著提升吞吐。本文以 BingBertSQuAD 微调与 BERT 预训练两个官方示例为主线完整讲解 1-bit Adam 的算法原理、NCCL/MPI 双后端选型、freeze_step/cuda_aware/comm_backend_name等核心配置参数、动量掩码momentum mask机制与检查点checkpoint注意事项并结合仓库内源码与单元测试说明其底层实现与验证方式帮助读者在自己的训练任务中直接落地这套压缩通信方案。1. 1-bit Adam 是什么通信瓶颈下的压缩优化器在大规模分布式训练中数据并行data parallelism要求在每轮迭代后对梯度进行全局 AllReduce 同步。随着模型规模与 GPU 数量增长通信开销逐渐成为主要瓶颈尤其是模型参数量越大、单次通信的字节数越多训练吞吐越受制于网络带宽。1-bit Adam 的核心思想是在梯度通信阶段不再传输完整的 32 位或 16 位浮点梯度而是将其量化为 1 比特每个数值仅用符号位表示后再进行 AllReduce从而使通信体积最多减少 5 倍。配合错误补偿compression error compensation机制在压缩阶段维持与标准 Adam 相当的收敛速度。该算法最初于 2020 年随 DeepSpeed 发布其系统实现细节、理论分析见论文《1-bit Adam: Communication Efficient Large-Scale Training with Adams Convergence Speed》。需要先说明的是DeepSpeed 文档在该教程中明确提示2022 年 3 月发布的 0/1 AdamZeroOneAdam部分沿用了 1-bit Adam 的设计在 BERT、GPT-2、ImageNet 等任务上拥有更好的通信效率且与基线 Adam 收敛质量一致因此新任务建议优先尝试 0/1 Adam教程见 zero-one-adam.md若 0/1 Adam 无法达到基线收敛再回退到 1-bit Adam。1.1 1-bit Adam v2 的改进本教程对应的 1-bit Adam v22021/03/04 更新相对 v1 有 3 点关键变化新增 NCCL 后端实现基于 PyTorch 分布式自带的 NCCL 后端实现压缩通信相比 MPI 版集成度更好、使用更简单在以太网环境性能更优在 InfiniBand 环境与 MPI 版性能持平新增动量掩码momentum mask支持针对训练过程中梯度恒为零的参数避免压缩误差在动量中持续累积修复若干 bug。1.2 使用前提与限制官方文档给出了 4 条需要特别注意的约束NCCL 版要求 PyTorch 1.8当使用 64 个及以上 GPU 时需要 NCCL 2.8.3详见后文环境准备1-bit Adam 虽兼容 FP16 与 FP32但目前仅在混合精度/FP16 训练下验证过收敛性MPI 版实现目前与流水线并行pipeline parallelism不兼容频繁加载检查点可能损害 1-bit Adam 的收敛性原因见后文检查点与误差重置。2. 环境准备与双后端选型2.1 安装 DeepSpeed若尚未克隆 DeepSpeed 仓库可克隆并初始化其示例子模块BingBertSQuAD 与 BERT 预训练示例均在其中git clone https://github.com/microsoft/DeepSpeed cd DeepSpeed git submodule update --init --recursive cd DeepSpeedExamples/2.2 后端一NCCL 版推荐v2 新增NCCL 版利用 PyTorch 分布式中的 NCCL 后端完成压缩通信与 PyTorch 分布式集成度高、易用性好官方推荐优先选用。环境要求PyTorch 1.8GPU 数量 64 时需要 NCCL 2.8.3 以避免特定的 NCCL 运行时 bug。在文档编写时2021/03NCCL 2.8.3 尚未被 PyTorch 官方支持官方给出的解决方案是通过LD_PRELOAD注入# 1) 安装 NCCL 2.8.3在 CUDA 11 系统上验证可行 apt-get install -y libnccl22.8.3-1cuda11.0 libnccl-dev2.8.3-1cuda11.0 # 2) 将 LD_PRELOAD 指向库路径 LD_PRELOAD/usr/lib/x86_64-linux-gnu/libnccl.so.2.8.3可通过NCCL_DEBUGINFO查看 NCCL 日志中实际加载的版本若显示NCCL version 2.8.3cuda11.0即说明LD_PRELOAD生效。需要说明的是上述版本号与安装命令是文档撰写时的历史事实实际部署时应以当前 PyTorch/NCCL 官方支持情况为准。2.3 后端二MPI 版MPI 版依赖消息传递接口MPI实现高级通信原语。DeepSpeed 的 Docker 镜像中已打包相关依赖若使用其他构建环境需自行安装 MPI 与 mpi4pypip install deepspeed[1bit_adam]官方在 CUDA-Aware MPIMVAPICH2-GDR上做过验证理论上任何支持 CUDA-Aware 的通信库如带 CUDA-Aware 支持的 OpenMPI均可正常工作。使用deepspeedlauncher 启动 MPI 版训练时必须指定 launcher 类型deepspeed --launcher[mvapich|openmpi] script.py也可以使用标准的mpirunlaunchermpirun -np [#processes] -ppn [#GPUs on each node] -hostfile [hostfile] [MPI flags] python [training_script.py]3. 配置 1-bit Adam三个核心参数1-bit Adam 通过 DeepSpeed 的 JSON 配置文件启用只需把优化器类型设为OneBitAdam并增加 3 个专属参数。最简配置示例如下{ train_batch_size: 4096, train_micro_batch_size_per_gpu: 16, optimizer: { type: OneBitAdam, params: { lr: 4e-4, freeze_step: 23000, cuda_aware: false, comm_backend_name: nccl } }, fp16: { enabled: true } }3.1 freeze_step预热步数决定何时开启压缩freeze_step是压缩通信生效前的预热warmup步数。在预热阶段1-bit Adam 按照标准 Adam 的方式执行完整精度通信让 Adam 的方差项二阶矩充分估计预热结束后才切换到 1 比特压缩通信。官方给出如下取值策略经验策略将该任务总训练步数的 15%~25% 设置为预热步数这与 Adam 的方差/二阶矩项有关理论分析见 1-bit Adam 论文若效果符合预期可系统性地逐步减小该值以榨取更多性能官方示例脚本中的freeze_step已按各自任务调优到最佳值。3.2 cuda_aware仅用于 MPI 后端的 CUDA-Aware 开关cuda_aware用于指示底层 MPI 库是否支持 CUDA-Aware 通信因此只在 MPI 后端下有意义仅支持 InfiniBand 互联 CUDA-Aware MPI 库如 MVAPICH2-GDR 或带 CUDA-Aware 支持的 OpenMPI的环境置为false时可在以太网上训练但通信前后需要在 CPU 与 GPU 缓冲区之间进行发送端和接收端的显式内存拷贝性能有所损失使用 NCCL 后端时无需设置cuda_aware。3.3 comm_backend_name选择通信后端v2 新增comm_backend_name用于选择压缩通信的底层实现nccl使用 NCCL 版实现推荐要求 PyTorch 1.8mpi使用 MPI 版实现需配合cuda_aware设置。3.4 官方参数总览参数说明默认值freeze_step压缩通信生效前的预热步数100000cuda_aware指示底层 MPI 库是否支持 CUDA-Aware 通信仅 MPI 后端falsecomm_backend_name选择 NCCL 或 MPI 后端实现nccl以上默认值可见于 config-json.md 中的 1-bit Adam 参数表与源码 adam.py 中构造函数的默认参数freeze_step100000、cuda_awareFalse、comm_backend_namenccl一致。除上述 3 个专属参数外1-bit Adam 同样支持标准 Adam 的参数包括lr、betas、eps、weight_decay、bias_correction等示例见 config-json.md 中的完整 OneBitAdam 配置。3.5 (v2 新增) 动量掩码处理梯度恒为零的参数背景问题1 比特压缩无法精确表示数值零。如果某个参数在训练中梯度恒为零其动量exp_avg中就会持续累积压缩误差进而损害训练。典型场景BERT 预训练seq length 128时bert.embeddings.position_embeddings.weight的第 129~512 行梯度与动量恒为零——因为模型最多支持 512 的序列长度而训练只学到 128。解决方案在参数分组中通过exp_avg_mask指定这些参数。掩码与参数同形状值为 1 的位置保留动量值为 0 的位置强制归零从而避免压缩误差累积。示例如下将参数分组传入deepspeed.initializeoptimizer_grouped_parameters [ { params: [some_param], weight_decay: 0.01, exp_avg_mask: mask, # 与参数同形状的 0/1 张量 }, { params: [other_param], weight_decay: 0.01, }, ] model, optimizer, _, _ deepspeed.initialize( configconfig_dict, modelmodel, model_parametersoptimizer_grouped_parameters, )重要注意检查点中保存的掩码不会被使用因为掩码在训练过程中可能变化例如 BERT 的 seqlen 128 与 512 需要不同的掩码。因此每次训练都必须重新在训练脚本中提供exp_avg_mask。从源码 adam.py 可以看到压缩阶段对带掩码的分组执行exp_avg.mul_(group[exp_avg_mask])在 AllReduce 之后直接将掩码中为 0 位置的动量清零load_state_dict也专门保留了用户传入的exp_avg_mask并覆盖检查点中的旧值。4. 算法与底层实现压缩通信是如何工作的1-bit Adam 的训练过程分为两个阶段源码 adam.py 中的step()方法清晰地体现了这一流程阶段一预热阶段adam_freeze_keyFalse按标准 Adam 计算一阶矩exp_avg与二阶矩exp_avg_sq梯度以完整精度通信。当state[step] self.freeze_step时打印 OnebitAdam - starting compressed communication切换adam_freeze_keyTrue并关闭 DeepSpeed 引擎中的enable_backward_allreduce流水线场景下为pipeline_enable_backward_allreduce因为后续梯度不再走常规 AllReduce。阶段二压缩阶段adam_freeze_keyTrue先按标准 Adam 更新一阶矩exp_avg.mul_(beta1).add_(1 - beta1, grad)通过self.comm_backend_handle.compressed_allreduce(exp_avg, state[worker_error], state[server_error], ...)对一阶矩做 1 比特压缩 AllReduce——这正是通信量降低 5 倍的关键路径若分组带exp_avg_mask则将掩码中为 0 位置的动量清零之后以冻结后的二阶矩继续计算更新量update exp_avg / (exp_avg_sq.sqrt() eps)并施加权重衰减与学习率。压缩通信的底层通过抽象的后端句柄完成comm_backend_name nccl时实例化 NcclBackend并在源码中断言 PyTorch 1.8否则提示改用 MPI 后端mpi时实例化MpiBackend(cuda_aware)。错误补偿机制为了在压缩阶段维持收敛速度每个 GPU 维护 worker 侧的worker_error与 server 侧的server_error两个误差缓冲区本轮压缩时被丢弃的量化残差会累加进误差并在下一轮压缩前补偿进待压缩的动量从而保证压缩误差不会无限累积。压缩报文中每个数值只占 1 比特这正是最多 5 倍通信量缩减的直接来源。从仓库单元测试 test_onebit.py 可以看到该实现已覆盖以下验证场景FP32/FP16 双精度下的基础收敛测试、exp_avg_mask动量掩码正确性断言、检查点保存/加载后掩码保持与 worker/server 误差被重置的断言以及不同流水线拓扑num_pp1/2/4下的 FP16 流水线训练测试。5. 案例一BingBertSQuAD 微调5.1 数据与模型准备下载 SQuAD 数据集训练集train-v1.1.json、验证集dev-v1.1.json下载 HuggingFace 预训练 checkpoint 与配置文件bert-large-uncased-whole-word-masking的模型权重与 json 配置也可以使用来自 DeepSpeed、HuggingFace 或 TensorFlow 的任意预训练 BERT 模型 checkpoint 进行微调。关于 checkpoint 加载、参数解析、初始化、前向/反向、权重更新与评估的细节可参考仓库中的 bert-finetuning.md 教程。5.2 运行方式官方提供了 3 组启动脚本分别对应NCCL 版、以太网上的 MPI 版、InfiniBand 上的 MPI 版其中 MPI 版又同时提供deepspeedlauncher 与mpirun两种启动方式。使用deepspeedlauncher无需手动设置 mpirun 参数节点数与 GPU 数自动检测bash run_squad_deepspeed_onebitadam.sh PATH_TO_OUTPUT_DIR使用mpirunmpirun -np [#processes] -ppn [#GPUs on each node] -hostfile [hostfile] [MPI flags] bash run_squad_mpi_onebitadam.sh例如使用 32 张 GPU8 节点 × 4 GPU/节点且启用 InfiniBand 时可借助 MVAPICH2 的 mpirun 启动mpirun -np 32 -ppn 4 -hostfile hosts -env MV2_USE_CUDA1 -env MV2_SUPPORT_DL1 -env MV2_ENABLE_AFFINITY0 -env MV2_SMP_USE_CMA0 bash run_squad_mpi_onebitadam.sh5.3 微调配置微调使用deepspeed_onebitadam_bsz96_config.json配置文件在运行训练脚本时通过--deepspeed启用 DeepSpeed并用--deepspeed_config指定该配置。官方实验使用的参数如下参数值总批大小Total batch size96每 GPU 微批大小Train micro batch size per GPU3优化器OnebitAdam学习率3e-5序列长度384权重衰减0.0训练轮数2freeze_step400comm_backend_namenccl5.4 效果与精度结果官方文档给出的准确率结果如下总批大小 96、32 张 GPU、2 轮训练学习率固定为 3e-5并尝试了多组种子与学习率后选择最优方案模型精度EMF1HuggingFace 基线bert-large-uncased-whole-word-maskingFP1687.2693.321-bit Adam 在 SQuAD 微调上取得了与 HuggingFace 基线持平或更优的 EM/F1 成绩。训练速度与扩展性scalability的详细数据见 1-bit Adam 的官方博客与论文。6. 案例二BERT 预训练6.1 数据准备与运行数据下载与预处理流程见 bert-pretraining.md 教程。官方在bing_bert/1-bit_adam/下同样提供 NCCL 版、以太网 MPI 版、InfiniBand MPI 版共 3 组脚本。使用deepspeedlauncherbash ds_train_bert_onebit_bsz4k_seq128.sh使用mpirunmpirun -np [#processes] -ppn [#GPUs on each node] -hostfile [hostfile] [MPI flags] bash mpi_train_bert_onebit_bsz4k_seq128.sh例如 32 张 GPU4 GPU/节点 × 8 节点配合 InfiniBand 与 MVAPICH2mpirun -np 32 -ppn 4 -hostfile hosts -env MV2_USE_CUDA1 -env MV2_SUPPORT_DL1 -env MV2_ENABLE_AFFINITY0 -env MV2_SMP_USE_CMA0 bash ds_train_bert_onebit_bsz4k_seq128.sh6.2 预训练配置BERT-largeseqlen 128官方提供的deepspeed_bsz4k_onebit_config_seq128_*.json配置如下{ train_batch_size: 4096, train_micro_batch_size_per_gpu: 16, steps_per_print: 100, prescale_gradients: false, optimizer: { type: OneBitAdam, params: { lr: 4e-4, weight_decay: 0.01, bias_correction: false, freeze_step: 23000, comm_backend_name: nccl } }, gradient_clipping: 1.0, fp16: { enabled: true, loss_scale: 0, initial_scale_power: 16 } }官方同时给出其他规模的建议取值BERT-baseseqlen 128建议freeze_step为 16000seqlen 512 预训练BERT-base 与 BERT-large 均建议freeze_step为 1500无论哪种规模都要按前文说明正确设置comm_backend_name与cuda_aware。6.3 预训练性能BERT 预训练下的通信压缩收益与加速比数据详见 1-bit Adam 官方博客与论文本文不再转述具体数字建议读者直接阅读原始资料获取实验细节。7. 检查点Checkpoint与压缩误差重置核心警告避免频繁加载检查点。1-bit Adam 依靠压缩误差补偿机制在压缩阶段维持收敛速度。加载检查点时代码会主动重置清零压缩误差官方给出了 3 条原因误差逐 GPU 不同每个 GPU 上的 worker/server 误差各不相同当前实现只在检查点中保存 rank 0 的误差因此必须重置。若想正确保存全部误差需要 O(num_gpu × model_size) 的内存来汇总内存开销极大理论上可以分布式保存但会使保存/加载逻辑复杂化GPU 数量必须一致即使能正确保存误差加载时也必须使用完全相同的 GPU 数量收敛影响有限官方在 BERT 预训练上验证过加载检查点时偶尔重置压缩误差不影响收敛。因此建议不要频繁加载检查点否则可能破坏误差补偿机制、影响收敛。源码层面adam.py 中的load_state_dict()实现了上述逻辑加载后若当前 step 小于freeze_step则保持/回到预热阶段恢复enable_backward_allreduce否则进入压缩阶段随后从每个参数的状态中弹出popworker_error与server_error完成误差重置。单元测试 test_onebit.py 中的TestOneBitAdamCheckpointing也专门断言了加载检查点后worker_error/server_error不存在以及exp_avg_mask在保存/加载前后保持不变。8. 从 1-bit Adam 到 0/1 Adam后续演进1-bit Adam 的经验被 0/1 AdamZeroOneAdam继承并进一步优化0/1 Adam 通过**自适应方差冻结adaptive variance freezing**与对优化器状态的 1 比特同步实现了更好的通信效率并在 BERT、GPT-2、ImageNet 等任务上保持与基线 Adam 相同的最终模型质量。配置方式与 1-bit Adam 同构但专属参数变为var_freeze_step、var_update_scaler、local_step_scaler、local_step_clipper等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 } }完整的参数说明见 config-json.md专项教程见 zero-one-adam.md。另外DeepSpeed 还提供了面向 LAMB 优化器的 1-bit 变体 OneBitLamb配置示例同样见 config-json.md教程见 onebit-lamb.md核心压缩思路一致。9. 总结与实践清单将 1-bit Adam 落地到自己的训练任务可按以下清单操作确认环境PyTorch 1.8NCCL 后端64 GPU 时确认 NCCL 2.8.3必要时用LD_PRELOAD注入若用 MPI 后端则安装deepspeed[1bit_adam]并准备 CUDA-Aware MPI 库开启 FP16 混合精度fp16.enabled true1-bit Adam 目前仅在 FP16 下验证收敛配置优化器type: OneBitAdam按任务规模设置freeze_step建议先取总训练步数的 15%~25%再逐步下调comm_backend_name选nccl推荐或mpi配合cuda_aware提供动量掩码对梯度恒为零的参数如 BERT 位置嵌入的越界行在参数分组中传入exp_avg_mask并在每次训练时都重新提供注意流水线限制MPI 版与流水线并行不兼容NCCL 版在压缩阶段会自动关闭 backward AllReduce谨慎处理检查点避免频繁加载检查点加载后压缩误差会被重置可能影响收敛。按此流程即可在通信受限的集群上以 1-bit Adam 获得最多 5 倍的通信量缩减与最高 3.4 倍的端到端训练加速具体收益取决于网络环境与任务规模官方博客与论文提供了完整的性能评估。赞分享推理引擎大模型【免费下载链接】FlexGenRunning large language models on a single GPU for throughput-oriented scenarios.项目地址https://gitcode.com/gh_mirrors/fl/FlexGen点击查看免费下载相关推荐DeepSpeed 1-bit Adam 实战指南通过最高 5 倍通信压缩加速分布式训练DeepSpeed 1 bit Adam 实战指南通过最高 5 倍通信压缩加速分布式训练 本篇技术指南围绕 DeepSpeed 中的 1 bit Adam 优人工智能大模型深度学习分布式训练预训练强化学习模型优化DeepSpeed 1-bit Adam 深度解析以 5 倍通信量压缩实现分布式大模型高效训练DeepSpeed 1 bit Adam 深度解析以 5 倍通信量压缩实现分布式大模型高效训练 1 bit Adam 是 DeepSpeed 中面向通信受限集推理引擎大模型革命性分布式优化1-bit Adam加速大模型训练5倍革命性分布式优化1 bit Adam加速大模型训练5倍 还在为大模型训练时的通信瓶颈和内存占用而烦恼DeepSpeed的1 bit Adam优化器让你用1/示例工程上一篇智能歌词管理革命163MusicLyrics 让音乐学习与收藏更高效下一篇如何在Audacity中免费获得专业级AI音频处理能力OpenVINO插件完整指南创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考