deepspeed zero3 llamafactory 保存checkpoint后第一step 就 OOM4张16g显卡 训练14b模型{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}}# # 模型与路径配置# model_name_or_path: /public/home/a15657581978/merged_correct_2epoch# 对话模板格式alpaca格式# 其他可选llama3, chatml, qwen, vicuna等template: alpaca# # 数据集配置# dataset_dir: /public/home/a15657581978/data-new/binding_sft_real_paired_512/short_tokens dataset: binding_part_000 cutoff_len:1024# 数据加载优化 streaming:falsepreprocessing_num_workers:4overwrite_cache:false# 内存防爆核心 dataloader_num_workers:2dataloader_pin_memory:falsedataloader_prefetch_factor:2# # 保存 checkpoint 时不保存 global_step# 防止内存溢出 oom# save_only_model:true# # 训练阶段与微调方法# stage: sft do_train:truefinetuning_type: lora lora_rank:64lora_alpha:128lora_dropout:0.05lora_target: q_proj,v_proj,k_proj,o_proj,up_proj,gate_proj,down_proj# # 海光 DCU 核心防崩溃配置关键# # additional_target: embed_tokens,lm_headflash_attn: disabled# attn_implementation: sdpaddp_backend: nccl pure_bf16:falsebf16:truefp16:falselow_cpu_mem_usage:truegradient_checkpointing:trueuse_cache:falseddp_find_unused_parameters:false# # 批次与训练参数# per_device_train_batch_size:1per_device_eval_batch_size:1gradient_accumulation_steps:16learning_rate:0.00005num_train_epochs:-1max_steps:18750warmup_ratio:0.1lr_scheduler_type: cosine# # 输出与日志# output_dir: /public/home/a15657581978/output/binding_part_000_512 logging_steps:10save_steps:100save_total_limit:2max_samples:1400000# # 分布式与监控# ddp_timeout:180000report_to:[tensorboard]logging_dir: /public/home/a15657581978/tblogs deepspeed: /public/home/a15657581978/training_configs_512/deepspeed_z3_offload_4x16g.json关键发现OOM 不是在保存中而是在保存后继续训练的第 101 步看你之前的日志时间线08:17:53 ✅ Checkpoint global_step100 is ready now! 保存成功完成 08:17:57 Rank 1 OOM in compute_loss → forward 4秒后第101步前向时崩保存本身是成功的OOM 发生在保存完成之后、训练恢复的那个 forward 里。那为什么第 1~100 步都不 OOM偏偏第 101 步 OOM因为保存过程搞脏了GPU 显存。具体来说有 3 个机制机制 1ZeRO-3 的参数缓存被清空了最关键ZeRO-3 的参数不是常驻 GPU 的而是用缓存池管理正常训练step 1~100的稳态 ┌─────────────────────────────┐ │ GPU 显存 │ │ ┌─────────┐ ┌─────────┐ │ │ │当前层参数│ │下一层参数│ ← 缓存预热好了 │ │ └─────────┘ └─────────┘ │ │ 剩余 ~1.5GiB 给临时张量 │ ← softmax 36MiB 够用 └─────────────────────────────┘ 保存 checkpoint 时 DeepSpeed 需要把参数从 GPU 拷到 CPU 再写磁盘 → 参数缓存池被腾空/打乱 保存完后恢复训练step 101 ┌─────────────────────────────┐ │ GPU 显存 │ │ ┌───────┐ ┌───────┐ ┌────┐ │ │ │碎片参数│ │碎片参数│ │碎片│ │ ← 缓存是冷的 │ │ └───────┘ └───────┘ └────┘ │ │ 14.81GiB 已分配, 571MiB 碎片 │ ← softmax 36MiB 分配不出来 │ 0 bytes 真正空闲 │ └─────────────────────────────┘你的报错信息也印证了这一点14.81 GiB is allocated by PyTorch, and571.62 MiB is reserved by PyTorch but unallocated571MiB 的显存被 PyTorch 占着但没法用就是碎片。机制 2保存时的序列化过程留下临时张量DeepSpeed 保存时会做这些事把 GPU 上的参数张量.cpu()拷到内存用torch.save()序列化写磁盘释放临时张量但在步骤 1→2→3 之间GPU 和 CPU 之间有大量数据搬运PyTorch 的内存分配器会留下碎片和未释放的临时缓冲区。在 NVIDIA CUDA 上这通常能自动回收但在海光 DCUHIP上内存回收机制可能没那么及时。机制 3保存后 ZeRO-3 AllGather 的冷启动峰值第 101 步的 forward 需要通过 AllGather 从其他卡收集当前层的完整参数。在稳态训练时这个 AllGather 是和上一层的计算重叠的峰值显存可控。但保存后第一步是冷启动稳态训练时step 99→100 层N计算的同时层N1的参数已经在后台AllGather好了 → 峰值显存 计算层N 缓存层N1 ≈ 可控 冷启动时step 100→101保存后 缓存是空的层1需要现场AllGather → 峰值显存 AllGather临时缓冲区 层1参数 优化器碎片 → 超了根本原因总结你的 4×16G 卡跑 Qwen2 ZeRO-3 LoRA 正常训练时显存占用 ≈ 14.5 GiB / 16 GiB只剩 1.5GiB 余量 → 1.5GiB 刚好够 softmax 等临时操作 保存 checkpoint 后 内存碎片 缓存清空 冷启动峰值 → 实际可用连续空间 36 MiB → softmax 分配 36 MiB 失败 → OOM不是保存本身吃显存而是保存搅乱了显存布局导致恢复训练时那 1.5GiB 的救命余量没了。怎么解决第 1 步开 expandable_segments解决碎片问题最可能一招见效exportPYTORCH_HIP_ALLOC_CONFexpandable_segments:True,max_split_size_mb:128这会让 PyTorch 的显存分配器支持动态扩展内存段571MiB 的碎片就能被重新利用36MiB 的 softmax 就能分配成功。第 2 步减小 DeepSpeed 通信桶降低峰值显存{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 步降低序列长度如果业务允许cutoff_len:768# 从 1024 降到 768直接砍掉 ~25% 的激活显存第 4 步拉大保存间隔减少被搅乱的频率save_steps:500# 从 100 改成 500减少保存次数第 5 步终极方案加 CPU Offload如果上面都不够让优化器状态住 CPU彻底释放 GPU 显存{zero_optimization:{stage:3,offload_optimizer:{device:cpu,pin_memory:true},offload_param:{device:cpu,pin_memory:true},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}}CPU Offload 会让训练慢 20%~40%但能省出 3~5GiB GPU 显存OOM 问题基本就不会再出现了。完整启动脚本#!/bin/bashpkill-9python;pkill-9torchrun;sleep3# 显存救命环境变量 exportPYTORCH_HIP_ALLOC_CONFexpandable_segments:True,max_split_size_mb:128# 指定4张卡 exportROCR_VISIBLE_DEVICES0,1,2,3# 超时时间拉长 exportNCCL_TIMEOUT1800# 启动训练 FORCE_TORCHRUN1llamafactory-cli train\/public/home/a15657581978/training_configs_512/binding_paired_000.yaml先只加PYTORCH_HIP_ALLOC_CONF这一个变量试试大概率就能从每次保存后 OOM变成能稳定跑完。如果还不行再加上第 2~5 步的修改。