海光 异构卡2 DCU (4×16G) 跑 DeepSpeed ZeRO-3 完整避坑指南

📅 2026/6/16 15:35:06
海光 异构卡2 DCU (4×16G) 跑 DeepSpeed ZeRO-3 完整避坑指南
海光 异构卡2 DCU (4×16G) 跑 DeepSpeed ZeRO-3 完整避坑指南一、 核心问题总结为什么总是崩溃pip show deepspeed|grepVersion Version:0.14.2das.opt1.dtk25041deepspeed 0.14.2das.opt1.dtk25041原来的配置{train_batch_size:auto,train_micro_batch_size_per_gpu:auto,gradient_accumulation_steps:auto,gradient_clipping:auto,zero_allow_untested_optimizer:true,fp16:{enabled:auto,loss_scale:0,loss_scale_window:1000,initial_scale_power:16,hysteresis:2,min_loss_scale:1},bf16:{enabled:auto},zero_optimization:{stage:3,overlap_comm:true,contiguous_gradients:true,sub_group_size:1e9,reduce_bucket_size:1e8,stage3_prefetch_bucket_size:1e8,stage3_param_persistence_threshold:1e6,stage3_max_live_parameters:3e8,stage3_max_reuse_distance:3e8,stage3_gather_16bit_weights_on_model_save:false}}你在训练中遇到了三个大坑根本原因如下保存后 OOM第 101 步崩溃不是保存本身吃显存而是保存过程破坏了 GPU 显存的缓存布局产生了大量碎片571MB 碎片却分配不出 36MB。在 DCU/HIP 环境下必须开启显存动态扩展机制。保存时死锁/超时YAML 中设置了save_only_model: true这在 ZeRO-3 中是致命的。ZeRO-3 必须保存优化器状态才能维持分片结构跳过优化器会导致多卡保存流程分裂互相等待直到超时。DeepSpeed 0.14.2 的 Bug该版本的 ZeRO-3 与框架的梯度累积no_sync机制冲突。必须在 ds_config 中强行关闭 DeepSpeed 层面的累积交由 HuggingFace Trainer 处理。二、 最终配置文件直接复制使用1. 启动脚本 (run.sh)最关键的是加上PYTORCH_HIP_ALLOC_CONF这是解决保存后 OOM 的核心#!/bin/bash# 1. 清理残留进程 pkill-9pythonpkill-9torchrunpkill-9llamafactory-clisleep3# 2. 核心环境变量 # 【救命配置】解决显存碎片化导致的 36MiB OOMexportPYTORCH_HIP_ALLOC_CONFexpandable_segments:True,max_split_size_mb:128# 指定使用 4 张卡海光 DCU 用 ROCRNVIDIA 用 CUDAexportROCR_VISIBLE_DEVICES0,1,2,3# 防止多卡保存时因硬盘慢而通信超时30分钟exportNCCL_TIMEOUT1800# 3. 启动训练 FORCE_TORCHRUN1llamafactory-cli train\/public/home/a15657581978/training_configs_512/binding_paired_000.yaml2. DeepSpeed 配置 (deepspeed_z3_offload_4x16g.json)注意所有浮点数必须写成纯整型如1e9写成1000000000gradient_accumulation_steps必须写死为1{train_batch_size:auto,train_micro_batch_size_per_gpu:auto,gradient_accumulation_steps:1,gradient_clipping:auto,zero_allow_untested_optimizer:true,fp16:{enabled:auto,loss_scale:0,loss_scale_window:1000,initial_scale_power:16,hysteresis:2,min_loss_scale:1},bf16:{enabled:auto},zero_optimization:{stage:3,overlap_comm:false,contiguous_gradients:true,sub_group_size:1000000000,reduce_bucket_size:200000000,stage3_prefetch_bucket_size:20000000,stage3_param_persistence_threshold:1000000,stage3_max_live_parameters:500000000,stage3_max_reuse_distance:500000000,stage3_gather_16bit_weights_on_model_save:false}}3. YAML 训练配置关键修改项在你的 YAML 中必须确认以下三个参数的值其余保持不变# 【必须改】ZeRO-3 必须保存优化器否则保存时死锁save_only_model:false# 【必须改】加大超时时间防止硬盘写入慢导致多卡失联ddp_timeout:18000000# 【不用改】保持 16由 HuggingFace Trainer 做累积绕开 0.14.2 的 Buggradient_accumulation_steps:16# 【建议改】如果加了环境变量还是偶发 OOM把序列长度降低cutoff_len:768# 【建议改】减少保存频率降低保存后冷启动带来的 OOM 风险save_steps:500三、 训练结束后的模型导出因为你设置了stage3_gather_16bit_weights_on_model_save: false训练保存的 checkpoint 是零碎的分片不能直接用 transformers 加载。必须用 LLaMA-Factory 的导出命令在 CPU 上拼接并合并 LoRAllamafactory-cliexport\--model_name_or_path/public/home/a15657581978/merged_correct_2epoch\--adapter_name_or_path/public/home/a15657581978/output/binding_part_000_512\--templatealpaca\--finetuning_typelora\--export_dir/public/home/a15657581978/exported_final_model\--export_devicecpu\--export_legacy_formatfalse--export_device cpu是关键借用系统内存拼接完整模型避免 16G 显卡被撑爆。四、 终极降级方案如果按照上述配置配置后依然在第 101 步 OOM说明 4×16G 跑 ZeRO-3 确实顶不住 Qwen2 的显存波动。此时请降级到 ZeRO-2稳定优先修改ds_config.json{train_batch_size:auto,train_micro_batch_size_per_gpu:auto,gradient_accumulation_steps:auto,gradient_clipping:auto,zero_allow_untested_optimizer:true,fp16:{enabled:auto},bf16:{enabled:auto},zero_optimization:{stage:2,overlap_comm:false,contiguous_gradients:true,reduce_bucket_size:300000000,allgather_bucket_size:300000000}}ZeRO-2 不切分模型参数没有保存时的 AllGather 冷启动问题极其稳定但能支持的模型上限比 ZeRO-3 小。这个发现非常关键你亲眼看到了真相save_only_model: true在 ZeRO-3 下根本没用你设置了save_only_model: true本以为只存模型结果4 × model_states.pt (7.65GB × 4 30.6GB) 4 × optim_states.pt (825MB × 4 3.3GB) ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 总计: ~33.9GB 一个 checkpoint优化器状态一个都没少存为什么会这样save_only_model: true 的本意 HuggingFace Trainer → 我只保存模型权重优化器、调度器、随机种子全扔了 → 省空间 ✓ 但在 ZeRO-3 下实际发生的事 HuggingFace Trainer → 告诉 DeepSpeed 只存模型 DeepSpeed → 我不听ZeRO-3 的模型碎片和优化器是绑定的必须一起存 → 照样存了 model_states optim_states → save_only_model: true 形同虚设在 ZeRO-3 中优化器状态里藏着模型参数的 fp32 主副本和分片映射关系DeepSpeed 不可能只存模型不存优化器。所以你之前的注释# 保存 checkpoint 时不保存 global_step# 防止内存溢出 oomsave_only_model:true# ← 这个在 ZeRO-3 下根本不生效这个设置既没有省空间反而还可能制造混乱HuggingFace Trainer 以为优化器没存跳过了某些元数据的保存导致恢复训练时状态不对。这个发现也解释了 OOM 的真正原因写33.9GB的数据到硬盘耗时非常长。你看之前的日志08:17:40 开始保存 08:17:53 保存完成 ← 整整 13 秒4张卡各写了 ~8.5GB 08:17:57 第 101 步 OOM13 秒的磁盘 IO 期间4 张卡的 GPU 显存布局被彻底搅乱参数缓存池被清空内存碎片大量产生恢复训练时的 AllGather 冷启动峰值更高这才是每次保存后就 OOM的根本原因。解决方案第 1 步把save_only_model改成false既然 DeepSpeed 反正都要存优化器不如让 HuggingFace Trainer 和 DeepSpeed 的保存行为保持一致避免元数据缺失导致恢复训练时出问题save_only_model:false第 2 步大幅拉长保存间隔 限制保存数量33.9GB 一个 checkpoint太大了。减少保存频率降低保存后 OOM的风险save_steps:500# 从 100 改成 500减少被打断的次数save_total_limit:1# 只保留最新的 1 个防止硬盘被撑爆# 33.9GB × 2 67.8GB很快就满了第 3 步开启动态显存分配必须exportPYTORCH_HIP_ALLOC_CONFexpandable_segments:True,max_split_size_mb:128这是解决保存后碎片化导致 36MiB 分配不出来的关键。修改汇总只列出需要改的行YAMLsave_only_model:false# ZeRO-3 下这个没意义改成 false 保持一致性save_steps:500# 减少保存频率save_total_limit:1# 只留最新的 1 个ddp_timeout:18000000# 30 分钟超时启动脚本#!/bin/bashpkill-9python;pkill-9torchrun;sleep3exportPYTORCH_HIP_ALLOC_CONFexpandable_segments:True,max_split_size_mb:128exportROCR_VISIBLE_DEVICES0,1,2,3exportNCCL_TIMEOUT1800FORCE_TORCHRUN1llamafactory-cli train\/public/home/a15657581978/training_configs_512/binding_paired_000.yamlds_config.json不变保持之前改好的版本{gradient_accumulation_steps:1,overlap_comm:false,reduce_bucket_size:200000000,stage3_prefetch_bucket_size:20000000,stage3_max_live_parameters:500000000,stage3_max_reuse_distance:500000000,stage3_gather_16bit_weights_on_model_save:false}一张图总结你原来的理解 实际情况 ┌─────────────────────┐ ┌─────────────────────┐ │ save_only_model:true │ │ save_only_model:true │ │ → 只存模型 │ │ → DeepSpeed 无视 │ │ → 省空间 ✓ │ │ → 照存 33.9GB │ │ → 防 OOM ✓ │ │ → 13秒磁盘IO │ └─────────────────────┘ │ → 保存后碎片化 │ │ → 第101步 OOM │ └─────────────────────┘ 改后的策略 ┌─────────────────────────────────────────┐ │ save_only_model: false 保持一致性 │ │ save_steps: 500 减少被打断 │ │ save_total_limit: 1 省硬盘 │ │ PYTORCH_HIP_ALLOC_CONF 防碎片化 │ │ ds_config: 小桶 不重叠 降峰值显存 │ └─────────────────────────────────────────┘原来你写的是auto啊这就更危险了我来给你揭秘auto在背后到底干了什么你就明白为什么必须强行改成1了。auto的真相它最终还是变成了 16在 DeepSpeed 配置文件中当你写gradient_accumulation_steps: auto时它的意思是“DeepSpeed你别管了去问 Hugging Face Trainer 要这个数值。”然后发生的流程是这样的DeepSpeed 看到auto跑去问 HF Trainer。HF Trainer 看了你 YAML 里的配置gradient_accumulation_steps: 16。HF Trainer 告诉 DeepSpeed“是 16”DeepSpeed 底层引擎默默地把auto替换成了16并接管了梯度累积的逻辑。所以你写auto和写16在 DeepSpeed 眼里完全是一回事。为什么auto会触发 0.14.2 版本的死锁 Bug当 DeepSpeed 拿到gradient_accumulation_steps 16后它会尝试自己控制no_sync()即前 15 步不进行跨卡梯度同步第 16 步才同步。但在 DeepSpeed 0.14.2 和 Hugging Face Trainer 的结合中这俩哥们抢夺了控制权HF Trainer 也想发no_sync()指令。DeepSpeed 底层也在发no_sync()指令。结果就是状态机错乱。在保存 checkpoint 时DeepSpeed 以为还在no_sync期间不需要同步而其他卡以为已经同步了最终导致多卡互相等待死锁超时为什么必须强行写1把ds_config.json里的gradient_accumulation_steps强行写成1本质上是缴了 DeepSpeed 的械gradient_accumulation_steps:1这告诉 DeepSpeed“你每一步都给我老老实实做梯度同步不要搞什么累积延迟”那你的有效 Batch Size 会变吗完全不会因为 YAML 里的gradient_accumulation_steps: 16没变HF Trainer 依然会在外层控制连续算 16 步微批次把 loss 加起来然后再更新一次权重。区别仅仅是写auto(实际为16)由 DeepSpeed 底层控制跨卡通信延迟触发死锁 Bug。写1由 HF Trainer 控制每一步都正常通信完美避开死锁 Bug数学结果完全等价。总结你的auto恰恰是导致你之前保存时死锁/超时的元凶之一请务必在ds_config.json中把它改成整数1这是在 DeepSpeed ZeRO-3 下保证稳定性的金科玉律。deepspeed 已知的bug下面按「版本区间 问题类型」把 DeepSpeed尤其是 ZeRO-3 相关的典型坑整理一下重点是你踩过的那几个梯度累积 /no_sync、保存 checkpoint 死锁/假死、OOM 等。说明DeepSpeed 目前主线版本是 0.16.x / 0.19.x“2.61 以上”这个版本号在官方发行版里并不存在所以我按 0.10、0.14–0.16、0.19 这样的大版本来总结。一、和 ZeRO-3 / 梯度累积no_sync相关的典型 Bug1.no_sync与 ZeRO-2/3 梯度分区冲突你踩的核心坑典型报错AssertionError: no_sync context manager is incompatible with gradient partitioning logic of ZeRO stage 2/3本质原因ZeRO-2/3 会对梯度做分区每个 GPU 只存一部分需要严格保证梯度同步的时机PyTorch 的no_sync()上下文管理器在梯度累积时故意跳过 All-Reduce两者同时启用时梯度分区的同步逻辑和no_sync的“不同步”逻辑互相打架导致断言失败或死锁。这个问题在社区里多次被提到是 ZeRO-2/3 梯度累积的经典冲突。涉及版本DeepSpeed 0.14.x ~ 0.16.0 都有用户反馈此问题尤其在 0.16.0 上特别明显。社区普遍推荐降级到 0.15.4可以规避。0.15.4 的 release notes 主要是一些补丁修复并没有专门提到这个问题的根本修复只是该版本在实践上更稳定。工程上的通用规避方法不限版本在 ds_config 里把 DeepSpeed 侧的梯度累积关掉gradient_accumulation_steps:1然后在 HuggingFace Trainer / 自己训练循环里控制累积步数YAML 里gradient_accumulation_steps: 16不用改。这样 DeepSpeed 每一步都同步梯度不进入no_sync模式就不会和 ZeRO-2/3 的梯度分区冲突。尽量避免手动使用no_sync()ZeRO-2/3 本身已经有复杂的通信调度不要再套一层no_sync否则几乎必然踩坑。2.gradient_accumulation_steps: auto导致 DeepSpeed 接管累积并和 Trainer 冲突现象ds_config 写gradient_accumulation_steps: autoDeepSpeed 会去读 Trainer 的设置最终变成 16DeepSpeed 自己控制no_sync和 HuggingFace Trainer 的梯度累积逻辑重叠在保存 checkpoint 时出现各种假死 / 死锁 / NCCL 超时。建议在 ds_config 中永远写死整数不要写autogradient_accumulation_steps:1把真正的累积步数交给 Trainer / 外层训练循环控制。二、保存 Checkpoint 时的死锁 / 假死 / NCCL 超时1. ZeRO-3 save_only_model: true导致多卡保存失步现象和你之前的情况一模一样日志里只有 Rank 0 在保存 model_states / optim_states看不到其他 Rank 的保存日志保存完后继续训练时Rank 1 在compute_loss/ forward 中崩溃NCCL 超时。根因ZeRO-3 的模型参数和优化器状态是深度绑定的优化器里还藏着参数的 FP32 主副本save_only_model: true让 HuggingFace Trainer 跳过优化器保存DeepSpeed 却还在按“必须存优化器”的逻辑走导致多卡保存流程对不上互相等屏障超时。解决方案ZeRO-3 下必须save_only_model:false并适当加大超时ddp_timeout:18000000# 30 分钟2. 保存耗时太长触发 NCCL 超时 / 被误判为挂掉现象4×16G 卡保存一个大 checkpoint几十 GB要十几秒甚至更久其他卡等不及超过ddp_timeout默认 180s就报 NCCL 超时 / SIGTERM。解决方案把ddp_timeout调大比如 1800s 或 18000s尽量把 checkpoint 写到本地 SSD 而不是慢速网络盘减少保存频率save_steps: 500而不是 100。3. ZeRO-3 checkpoint 恢复时加载失败 / 状态不对常见问题直接用torch.load()加载zero_pp_rank_*.pt会丢失分片信息导致参数拼不回来HuggingFace Trainer 在 ZeRO-3 下直接model.load_state_dict()也会报错或状态缺失。正确做法用 DeepSpeed 提供的load_checkpoint/save_checkpoint或官方zero_to_fp32.py脚本先合并分片再加载LLaMA-Factory 的export命令其实就是在做类似的事情并支持--export_device cpu来避免 GPU OOM。三、保存后继续训练时的 OOM显存碎片 / 冷启动峰值1. 保存后第 101 步 OOM而不是保存时 OOM现象和你看到的完全一致保存 checkpoint 本身成功显存没爆4 秒后恢复训练的第 101 步在softmax/ attention 等小分配上 OOMtorch.OutOfMemoryError: HIP out of memory. Tried to allocate 36.00 MiB.原因保存时把参数从 GPU 拷到 CPU 再写磁盘会打乱 ZeRO-3 的参数缓存池清空“热参数”缓存PyTorch/HIP 的显存分配器在频繁分配/释放后产生大量碎片你看到的 571MiB reserved but unallocated 就是碎片第 101 步是 AllGather 冷启动需要临时重新收集完整参数峰值显存比稳态高刚好把那 1.5GiB 余量吃光连 36MiB 都分配不出来。解决方案你现在已经用的那一套开启 HIP/CUDA 的动态段分配exportPYTORCH_HIP_ALLOC_CONFexpandable_segments:True,max_split_size_mb:128把 ZeRO-3 的通信桶调小降低峰值显存reduce_bucket_size:200000000,stage3_prefetch_bucket_size:20000000,stage3_max_live_parameters:500000000必要时降cutoff_len比如 1024→768或增大save_steps。四、0.10.x vs 0.14–0.16 vs 0.19 的典型差异 Bug 分布1. 0.10.x 时代早期 ZeRO-3特点ZeRO-3 刚引入通信和内存管理相对粗糙社区反馈较多的是ZeRO-3 通信死锁尤其和梯度累积、多机组合时某些 CPU Offload 场景下 OOM / crash。建议新项目不建议再用 0.10.x除非你锁死了老环境如果在用优先关闭overlap_comm把reduce_bucket_size/allgather_bucket_size调小。2. 0.14–0.16你踩坑最集中的版本核心问题ZeRO-2/3 no_sync/ 梯度累积冲突0.14.x ~ 0.16.0 上都有用户报告no_sync incompatible with gradient partitioning错误0.16.0 尤其明显社区建议降级到 0.15.4。checkpoint 保存死锁 / 假死ZeRO-3 save_only_model: true 慢磁盘 → 多卡失步 → NCCL 超时0.14/0.15 对 ZeRO-3 的保存路径还不够稳健容易出现 Rank 0 写完、其他 Rank 没对齐的情况。HIP/DCU 上的显存碎片问题在海光 DCU 等环境下0.14–0.16 对expandable_segments等新特性支持不完善更容易碎片化 OOM后来 0.19.x 才专门加了 Zero3 defragment utility 等工具。实用建议如果你在 0.14–0.16 上训练 ZeRO-3ds_configgradient_accumulation_steps: 1overlap_comm: falseYAMLsave_only_model: falseddp_timeout: 18000000环境变量PYTORCH_HIP_ALLOC_CONFexpandable_segments:True,max_split_size_mb:128能升级的话推荐升级到 0.15.4 或 0.19.x而不是留在 0.14.x / 0.16.0。3. 0.19.x当前最新线主要改进修复了进程组关闭时卡死的问题训练结束 hang 住新增 Zero3 defragment utility专门处理 ZeRO-3 的显存碎片问题修复 ZeRO-3 forward crash、test_zf.py hang 等稳定性问题对 CPU Offload、FP16/BF16 优化器状态、torch.compile / torch.func 等兼容性更好。仍然要注意的ZeRO-3 no_sync的根本问题在 0.19.x 依然存在这是设计层面的冲突所以仍然建议gradient_accumulation_steps: 1把累积交给上层框架0.19.x 对 PyTorch 版本要求更高需要 2.11 等升级时要核对兼容性。五、其他版本相关的常见坑非 ZeRO-3 专属安装 / 编译问题Windows / 某些 CUDA 版本下pip install deepspeed经常遇到dskernel/async_io编译失败需要从源码编译或禁用某些 ops多 CUDA 版本共存时PyTorch 编译用的 CUDA 和系统 CUDA 不匹配导致CUDAMismatchException。CPU Offload 相关CPU 内存不足时ZeRO-3 Offload 会直接被 OOM Killer 杀掉表现为 exit code -9某些版本 CPUAdam 的 C 扩展编译失败导致AttributeError: DeepSpeedCPUAdam object has no attribute ds_opt_adam。检查点转换 / 加载ZeRO-3 的 checkpoint 是分片的直接用 HuggingFacefrom_pretrained加载会报各种KeyError/TypeError必须先用zero_to_fp32.py合并。六、给你一个「版本选择 配置模板」的实用总结版本选择新项目优先 0.19.x稳定性与碎片管理都更好如果在 0.16.x 上踩到no_sync冲突至少升级到 0.15.4 或直接上 0.19.x尽量避免在 0.14.x / 0.16.0 上长期跑生产级 ZeRO-3 训练。ZeRO-3 通用配置要点适用于 0.14–0.19{gradient_accumulation_steps:1,overlap_comm:false,contiguous_gradients:true,reduce_bucket_size:200000000,stage3_prefetch_bucket_size:20000000,stage3_max_live_parameters:500000000,stage3_max_reuse_distance:500000000,stage3_gather_16bit_weights_on_model_save:false}YAML 里save_only_model:falseddp_timeout:18000000save_steps:500save_total_limit:1环境变量HIP/CUDA 都建议开exportPYTORCH_HIP_ALLOC_CONFexpandable_segments:True,max_split_size_mb:128exportNCCL_TIMEOUT1800按这套组合大部分你遇到的“保存后 OOM / 死锁 / no_sync 冲突”都可以稳定规避。