DeepSpeed-Chat 2023-08-31 版本解读Llama/Llama-2 全流程支持、MixZ 与 ZeRO-Offload 深度解析【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed本文基于 DeepSpeed 官方博客 blogs/deepspeed-chat/ds-chat-release-8-31/README.md 展开系统梳理 2023 年 8 月 31 日 DeepSpeed-Chat 发布的三大主题Llama/Llama-2 模型的三阶段 RLHF 全流程支持、面向 LoRA 场景的 Mixed Precision ZeROMixZ与 ZeRO-Offload 两项效率优化以及配套的稳定性修复与工程化改进。读完本文你将理解这些优化在 DeepSpeed 运行时中的配置项与源码落点如 deepspeed/runtime/zero/config.py、deepspeed/ops/quantizer/quantizer.py并能据此在实际项目中启用相应特性。1. 发布概述三条主线DeepSpeed-Chat 是一个通用的 RLHFReinforcement Learning from Human Feedback训练系统框架基于 OpenAI InstructGPT 的训练策略目标是让 ChatGPT 风格模型的训练变得“容易、快速、廉价且可扩展”。2023-08-31 的这次发布围绕三个方向推进Llama / Llama-2 模型的训练系统支持将 Hybrid Engine、ZeRO 家族优化、LoRALow-Rank Adaptation等能力完整集成到三阶段 RLHF 流水线中。借助 Hybrid EngineLlama-2-7B 与 Llama-2-13B 的体验生成experience generation阶段最高加速7.1 倍。效率与可及性提升MixZMixed Precision ZeRO在 ZeRO/ZeRO 基础上扩展的优化策略集专为“LoRA RLHF”场景裁剪Llama-2-70B 在 128 块 V100 上训练吞吐最高提升3.3 倍ZeRO-Offload将优化器显存与计算从 GPU 卸载到主机 CPU最多用1/16 的 GPU训练同等大模型。稳定性与软件增强修复了 DeepSpeed-Chat 客户端代码与 DeepSpeed 运行时的多项训练稳定性/收敛问题全部使用 OPT 模型族做过端到端验证同时新增特征扫描characterization脚本、TensorBoard 等观测能力并为 Step 3 引入每夜运行的 PyTest CI 工作流。官方建议试用该版本时搭配 DeepSpeedv0.10.2发布版与当时 DeepSpeedExamples 仓库中的最新 DeepSpeed-Chat 代码DeepSpeed-Chat 应用代码位于独立的 DeepSpeedExamples 仓库不在本仓库内本文仅引用本仓库中的博客与运行时源码。2. Llama / Llama-2 全流程系统支持2.1 支持过程中遇到并解决的工程问题DeepSpeed-Chat 的训练框架在这之后对 Llama 与 Llama-2 提供了全部三个阶段的系统支持。官方博客明确说明这一过程遇到的是一类从“小运行时错误到复杂性能问题”的光谱问题其中两个核心障碍是Llama 模型架构偏离标准 Transformers Block与 DeepSpeed 的推理内核不兼容也与 Hybrid Engine 所采用的 DeepSpeed container policy模型算子划分策略不匹配ZeRO 家族优化与 Hybrid Engine 加速推理内核的交互需要大量修改覆盖 DeepSpeed-Chat 流水线与 DeepSpeed 运行时两侧的代码。这些问题的解决保证了 ZeRO 家族优化、Hybrid Engine 与 LoRA 在 Llama 系模型上可以正确组合使用。2.2 集成的关键优化清单优化在 Llama/Llama-2 上的状态DeepSpeed-Chat 集成完整接入基于 InstructGPT 策略的三阶段端到端 RLHF 框架Hybrid Engine支持全部 Llama-1 变体、Llama-2-7B、Llama-2-13B生成阶段显著加速ZeRO 与 ZeRO-Offload完整支持offload 可利用整机内存容量支撑更大模型MixZ通过冻结权重低精度驻留改善 Llama-2-70B 等大模型效率LoRA完整支持冻结原始权重、只学习低秩分解矩阵对大幅降低存储需求2.3 上手方式三阶段 Llama-2 脚本当时新增的 Llama/Llama-2 训练脚本按三阶段组织位于 DeepSpeedExamples 仓库的 DeepSpeed-Chat 应用目录下本仓库不包含这些脚本阶段脚本目录DeepSpeedExamples 中Step 1 监督微调applications/DeepSpeed-Chat/training/step1_supervised_finetuning/training_scripts/llama2Step 2 奖励模型微调applications/DeepSpeed-Chat/training/step2_reward_model_finetuning/training_scripts/llama2Step 3 RLHF 微调applications/DeepSpeed-Chat/training/step3_rlhf_finetuning/training_scripts/llama2官方博客附有一条重要免责声明Llama/Llama-2 支持的系统层面经过了充分测试但不保证训练收敛实际使用中可能需要调参才能达到收敛。2.4 Hybrid Engine 性能评估A100与不带 Hybrid Engine 的基线相比Llama-2-7B 生成加速7.1xLlama-2-13B 加速5.4x官方 Figure 1见文首图片。V100Llama-2-7B 加速4x16 块 V100Llama-2-13B 加速2.1x32 块 V100 对比 16 块 V100 基线。更大型的 Llama-2-70B 收益与 ZeRO-Offload 带来的资源缩减见下文第 3、4 节。3. MixZLoRA 场景下 3.3 倍吞吐3.1 原理从 ZeRO 到 ZeRO 再到 MixZMixZ官方 PR #3954是构建在 ZeRO 与 ZeRO 之上的扩展优化策略集针对“RLHF LoRA”这种大量参数冻结、只训练少量低秩矩阵的场景裁剪。其三个层次的设计参数切分同 ZeRO把模型参数分区到各 GPU 上以减小显存占用只在需要时按需 gather分层切分 量化通信同 ZeRO分层切分让参数在可能的情况下全部驻留单节点内使 all-gather 通信发生在节点内带宽远高于跨节点权重在 gather 前先量化进一步压缩通信量混合精度驻留MixZ 的关键差异ZeRO 中参数始终以 fp16/bf16 存储、通信前后做量化/反量化而 MixZ 允许把冻结权重即 LoRA 中不训练的部分持久化存放在更低精度中。这带来三个直接收益通信开销显著下降、消除运行时量化开销、从而支持更大 batch size 换取更高效率。本仓库的 MixZ 教程 将其拆成两个可独立启用的组件与博客描述一一对应qwZ量化权重对冻结权重做块量化block-based quantization降低显存与 all-gather 通信量。与 ZeRO 不同的是MixZ 的 qwZ保持冻结权重常驻量化态因此运行时没有量化开销hpZ分层切分通过数据重映射data remapping与重计算recomputation消除跨节点参数 all-gather 通信且同时作用于 backward 与 generation 两个 pass。3.2 在源码中的落点MixZ 的能力由一组 DeepSpeed ZeRO 配置字段承载在当前仓库的 deepspeed/runtime/zero/config.py 中可以确认{ zero_optimization: { stage: 3, zero_quantized_nontrainable_weights: true, zero_hpz_partition_size: 16 } }zero_quantized_nontrainable_weightsconfig.py 中定义默认False对不可训练的 ZeRO 参数做量化以节省显存与通信——这正是 MixZ 中“冻结权重低精度驻留”的实现开关区别于对可训练参数通信做量化的zero_quantized_weightszero_hpz_partition_size默认 1见 config.py分层切分的节点内分区数。教程明确要求多机训练时该值应设为每节点 GPU 数如每节点 8 卡则设为 8单机训练则不应设置。量化路径的执行链在源码中可以追溯到ZeRO-3 优化器deepspeed/runtime/zero/stage3.py接收zero_quantized_nontrainable_weights并透传给参数分区逻辑deepspeed/runtime/zero/parameter_offload.py、deepspeed/runtime/zero/partition_parameters.py底层量化/反量化由 deepspeed/ops/quantizer/quantizer.py 中的 CUDA 内核完成——博客中“由高度优化的 CUDA 内核驱动”的表述由此得到印证。3.3 训练脚本侧的一行改动DeepSpeed 引擎在初始化时如果直接传入 LoRA 模型可以自动识别冻结参数但常见做法是先初始化基座模型、之后再转换成 LoRA。此时需要显式通知引擎识别冻结参数教程给出的做法是在转换完成后调用model, optimizer, _, lr_scheduler deepspeed.initialize( modelmodel, optimizeroptimizer, argsargs, configds_config, lr_schedulerlr_scheduler, dist_init_requiredTrue) # ... (custom code to convert base model to LoRA model) ... # 再次告知 DeepSpeed 引擎 LoRA 冻结参数以便持久化量化 model.optimizer.quantize_nontrainable_params()即官方博客所说“集成进训练脚本只需一行代码”的具体形态。3.4 性能与定位实验Llama-2-70B 的 RLHF Step 3 训练64 与 128 块 V100 两种配置。64 卡下 MixZ 相比 ZeRO-3 基线吞吐提升 2x128 卡下提升 3.3x官方 Figure 3见本节首图官方博客强调 MixZ 与 QLoRA 类方案的定位差异QLoRA 通过 NF4 数据类型、paged optimizer 等技巧追求单卡微调大模型MixZ 则是同样由量化驱动、但面向大规模多 GPU 扩展的方案且与 ZeRO-Offload、Hybrid Engine 等既有特性兼容测试规模达 Azure 上 384 块 GPU。另外从源码结构看当前仓库中 ZeRO 通信路径还进一步演进出 LoCo 机制zeropp_loco_param配置项含err_beta、reset_T参数见 config.py 与 stage3.py 中的all_to_all_loco_quant_reduce分支说明该方向在发布之后仍在持续演进。4. ZeRO-Offload16 倍 GPU 缩减4.1 原理ZeRO-Offload 利用“现代 GPU 集群中 CPU 内存通常是 GPU 显存的 2-3 倍”这一硬件事实把优化器状态显存与计算从 GPU 卸载到主机 CPU用整机内存容量支撑超大模型训练同时通过工程手段避免 CPU 低带宽成为瓶颈。官方博客给出的量化结论Figure 4见本节首图配置训练 Llama-2-7B 所需 V100 数DS-Chat ZeRO-316 块 LoRA4 块 ZeRO-Offload1 块HuggingFace 基线因显存不足无法运行即 ZeRO-Offload 最多可将所需 GPU 数降低16 倍。4.2 从“被禁用”到“全阶段启用”ZeRO-Offload 在 DeepSpeed-Chat 首发时因与 Hybrid Engine、LoRA 组合出现训练不稳定而被禁用。本次发布在 Hybrid Engine 与 LoRA 改进、并对 ZeRO Stage 2 / Stage 3 的多种特性组合做了充分测试后ZeRO-Offload 得以在 DeepSpeed-Chat 全部三个阶段启用。但博客同时保留了一条明确的限制ZeRO-Offload ZeRO Stage 2 开启 Hybrid Engine 且关闭 LoRA的组合目前不受支持因观察到训练不稳定。官方 Figure 5 展示了启用 ZeRO-Offload 后所有受支持配置在 16 块 V100 上的奖励分数收敛曲线使用 OPT-1.3B actor 与 OPT-350M critic 模型验证。4.3 在源码中的落点ZeRO-Offload 的配置项同样在 deepspeed/runtime/zero/config.py 中定义新版offload_optimizer优化器卸载与offload_param参数卸载两个字典配置取代了已弃用的cpu_offload/cpu_offload_param。Stage-3 优化器在初始化时接收cpu_offload、cpuadam_cores_perc等参数见 stage3.py其中cpuadam_cores_perc用于控制 CPU Adam 占用的核心比例——这正是博客所述“在不被 CPU 带宽瓶颈住的前提下利用主机内存”的运行时体现。5. 稳定性 Bug 修复官方在 16 块 V100、OPT-1.3B actor OPT-350M critic 的条件下对全部受支持配置做了扫描Figure 6见本节首图验证修复后的训练收敛。修复分为三个层面5.1 DeepSpeed-Chat 流水线修复训练稳定性KL 散度与序列生成调整DSE PR #620微调 PPO 训练器中 KL 散度以缩小新旧策略差异_generate_sequence()中移除generate()的min_length参数避免生成序列被人为拉长保留“生成崩塌”在训练严重不收敛时的自然显现修复compute_rewards()的 off-by-one 错误RLHF 训练函数在对话结束后将 reward 与 value 置零防止advantages与returns计算错误。LoRA 独立学习率DSE PR #633此前 LoRA 参数直接沿用整体学习率导致 ZeRO-2 系列配置训练不稳定。修复方式为在构造优化器参数组时把 LoRA 的lora_right_weight与lora_left_weight显式拆分为独立参数组并赋予 LoRA 专用学习率修复后收敛显著改善。本仓库中 LoRA 注入策略可参见 deepspeed/module_inject/policy.py。独立的 LoRA 学习率命令行参数DSE PR #685三个训练阶段均可单独指定 LoRA 学习率Step 3 中 Actor 与 Critic 各自独立。Bug 修复ZeRO-3 挂起DSE PR #636Step 3 中启用 ZeRO Stage 3 且world_size 1时观察到挂起。根因是各 rank 的 token 生成不同步——一个 rank 仍停留在_generate_sequence()另一个已进入actor_model()前向由于 DeepSpeed-Chat 流水线中模型配置对象的生命周期特性Transformers 库的synced_gpus自动检测未被触发。修复方式是在调用generate()时若使用 ZeRO Stage 3 则显式传入synced_gpusTrue。梯度检查点 仅 LoRA 参数优化DSE PR #658借助 Transformers 库的enable_input_require_grads工具函数为输入嵌入开启梯度使 Step 3 支持“梯度检查点 只优化 LoRA 参数”这一原本不支持的组合。另有外部贡献者修复的 argparse 问题DSE PR #576与未初始化硬编码参数传递问题DSE PR #584。5.2 Hybrid Engine 修复LoRA Fuse/UnfuseDS PR #3563OPT LoRA Hybrid Engine 的 Step 3 训练中出现 LoRA 权重张量尺寸不匹配——OPT container policy 未融合 LoRA QKV 权重而 Hybrid Engine 侧期望已融合。修复是在 Hybrid Engine 中同时提供 fused 与 unfused 两套 LoRA 路径本仓库中 Hybrid Engine 与 LoRA 策略的交互实现可参见 deepspeed/module_inject/containers/features/hybrid_engine.py 与 deepspeed/runtime/hybrid_engine.py。HE ZeRO-3 守卫DS PR #3883Hybrid Engine 在重置推理容器参数时正确检查 ZeRO Stage 3 是否启用并扩展相应单元测试。5.3 ZeRO Stage 3 修复GatheredParameters 竞态DS PR #3819GatheredParameters上下文中的竞态条件导致多种status: INFLIGHT错误通过移除 Hybrid Engine 重复传入的输入参数修复。训练/评估分离 InflightParamRegistryDS PR #3884ZeRO-3 的InflightParamRegistry改为训练与评估各用一份修复残留 in-flight 参数引发的错误。这些修复与 Hybrid Engine 相关修复一起使 ZeRO-Offload 得以进入 DeepSpeed-Chat 流水线。移除打印中的param.ds_tensorDS PR #3928小改动但它让执行得以推进到更深处帮助定位 Step 3 挂不同步点与 DSE PR #636 配合。6. 软件改进可观测性与可维护性6.1 特征扫描Characterization脚本DeepSpeed-Chat 的特性组合Hybrid Engine、ZeRO、LoRA 等非常多且相互作用复杂官方为此加入了系统扫描 Step 1/2/3 各种特性组合的 characterization 脚本默认 OPT可改为 Llama脚本分别位于 DeepSpeedExamples 中三个阶段的training_scripts/opt/single_node/sweep目录。以 Step 3 为例扫描矩阵为特性取值ZeRO Stage2, 3Hybrid EngineTrue, FalseZeRO-OffloadTrue, FalseLoRATrue, False运行方式DeepSpeedExamples/applications/DeepSpeed-Chat/training/step3_rlhf_finetuning$ \ bash training_scripts/opt/single_node/sweep/run_step3_sweep.sh每种特性组合的训练日志保存在形如z${z}_he_${he}_offload_${offload}_lora_${lora}的目录中便于后续对比分析。6.2 观测能力Instrumentation新增的观测参数贯穿三个阶段通过各阶段main.py的命令行参数开启参数说明适用阶段--print_loss每步打印 loss1--enable_tensorboard开启 Runtime Engine 层 TensorBoard 记录1, 2, 3开启训练流水线层 TensorBoard 记录3--tensorboard_pathTensorBoard 日志写入路径1, 2, 3--print_answers所有 rank 上打印 Actor 模型的 prompt 与回答3TensorBoard 在三个阶段均可启用Step 3 有额外细节每个阶段中enable_tensorboard开启的是Runtime Engine 层监控对应 DeepSpeed 配置 JSON 的 monitoring 模块反映在模型训练配置中为tensorboard: { enabled: enable_tensorboard, output_path: f{tb_path}/ds_tensorboard_logs/, job_name: f{tb_name}_tensorboard }Step 3同时初始化 Actor 与 Critic因此每个模型各有一个Runtime Engine 层 TensorBoard 监控此外还有一层位于模型引擎之上的Pipeline 层监控采集reward、actor_loss、actor_loss_sum、critic_loss、critic_loss_sum指标。6.3 测试Step 3 的每夜 CI为维持 DeepSpeed-Chat 框架的长期健康官方为Step 3 RLHF 训练新增了 PyTest 持续集成测试每夜运行。测试矩阵同样是 4 个二值参数构成的16 个配置ZeRO Stage 2/3 × HE 开/关 × ZeRO-Offload 开/关 × LoRA 开/关使用facebook/opt-125m作为 actor 与 critic。每个配置运行有限的 Step 3 非溢出训练步即 actor 与 critic 均不发生梯度溢出的步保存 actor/critic 模型后通过断言检查训练流水线执行正确性与模型保存完整性。7. 适用前提与实践建议版本对齐该博客描述的是 DeepSpeed v0.10.2 时期2023-08-31 发布的能力边界。本仓库当前版本中MixZ 相关的zero_quantized_nontrainable_weights、zero_hpz_partition_size与 LoCo 扩展等配置仍然存在且持续演进但 DeepSpeed-Chat 应用脚本已移至独立的 DeepSpeedExamples 仓库复现博客实验时应使用博客指定的仓库与提交版本而非本仓库中的脚本。特性组合限制ZeRO-Offload 虽已全阶段启用但“ZeRO-Offload ZeRO-2 HE 无 LoRA”组合不支持启用 MixZ 多机训练时zero_hpz_partition_size必须等于每节点 GPU 数单机则不设置。收敛预期管理官方博客明确不保证 Llama/Llama-2 的收敛性需要调参而 OPT 模型族经过了完整的端到端稳定性验证Figure 5/6 的扫描数据对稳定性敏感的场景可优先参考 OPT 的验证结论。深入阅读路径本仓库中与本文强相关的资料包括 DeepSpeed-Chat 主博客、MixZ 教程、ZeRO 教程、ZeRO 配置定义 deepspeed/runtime/zero/config.py 以及量化内核 deepspeed/ops/quantizer/quantizer.py可沿此路径从发布特性下钻到运行时实现。【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
