
DeepSpeed Zero-3 的参数分区与按需访问机制【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed大模型训练时用普通数据并行经常撞上torch.OutOfMemoryError: CUDA out of memory——每张卡都存一份完整参数内存随卡数线性增长加卡也解决不了单卡放不下。DeepSpeed 的 Zero-3 用参数分区解决它把参数切成 N 份每个数据并行进程只保留自己那份算到某一层时再从其他卡临时凑齐完整副本用完立刻释放。如何最快跑通 Zero-3克隆仓库git clone https://gitcode.com/GitHub_Trending/de/DeepSpeed并在 DeepSpeed 目录下执行pip install .。然后在训练脚本旁放一份最小配置{ zero_optimization: { stage: 3, contiguous_gradients: true, allgather_bucket_size: 500000000, reduce_bucket_size: 500000000, stage3_param_persistence_threshold: 100000, stage3_gather_16bit_weights_on_model_save: true }, bf16: {} }训练脚本里把优化器初始化换成deepspeed.initialize(modelmodel, model_parametersmodel.parameters(), configzero3.json)即可开训。stage3_*参数不用全写内置默认值在多数场景下可用等碰到内存墙或吞吐瓶颈再回头调。如何理解 ZeRO Stage 1/2/3 的分区对象与内存收益ZeRO 三个阶段是渐进的每一档都多分一类每卡冗余的数据。下图来自官方文档以 7.5B 参数模型 Adam 优化器 64 卡为例阶段分区对象上图示例中每卡内存适用区间基线数据并行不分区每卡全量120GB小模型Stage 1优化器状态如 Adam 的动量、方差31.4GB1B 级Stage 2优化器状态 梯度16.6GB1B~10B 级Stage 3参数 梯度 优化器状态1.9GB10B 级以上例子里每卡内存从 120GB 降到 1.9GB代价是通信量升到基线的约 1.5 倍。Stage 3 用异步预取把这部分通信藏进计算时间里——算当前层时下一层参数已经在路上。参数访问本身由一个状态机驱动每个参数处于三种状态之一定义见 partition_parameters.py 中的ZeroParamStatusNOT_AVAILABLE 本地只有分片不能直接计算 │ allgather 发起 ▼ INFLIGHT 分片正在网络上传输 │ 传输完成 ▼ AVAILABLE GPU 上有完整副本前向/反向可用 │ 使用完毕、引用计数归零 ▼ NOT_AVAILABLE 释放回分片形态内存还给别的参数围绕这个流转有两个优化点。其一是预取fetch_sub_module见 partitioned_param_coordinator.py在当前层还没算完时就发起后续层参数的收集参数到达时往往正好要用。其二是小参数驻留元素数低于 1e5默认值的参数不参与分区常驻每张卡省掉高频小参数反复 gather 释放的往返开销。如果某个参数被定义它的模块之外的代码引用典型如共享 embedding需要显式调用register_external_parameter登记依赖否则 DeepSpeed 会认为它没人用而提前释放。stage3_max_live_parameters 等相关参数怎么调这些配置都围绕一件事每卡驻留多少参数。常用的几项参数作用建议值调小的后果调大的后果stage3_max_live_parameters每卡释放前最多驻留的参数元素数1e9更省内存gather/释放更频繁更省通信、更吃内存可能 OOMstage3_max_reuse_distance距离内会被复用的参数不释放1e9更省内存密集复用的参数被反复释放更吃内存stage3_prefetch_bucket_size单次提前拉取的参数元素上限5e7预取不足计算等通信更吃内存重叠更好stage3_param_persistence_threshold元素数低于此值的参数不分区1e5小参数常驻变少延迟受限的小消息变多常驻小参数变多占更多显存allgather_bucket_size/reduce_bucket_sizeall-gather / 归约的桶大小5e8通信更碎、延迟更高通信缓冲更大stage3_gather_16bit_weights_on_model_save保存 16bit 权重时是否先聚齐全量权重truefalse 时 state_dict 不含完整权重—调参顺序建议全部保持默认 → OOM 就先降stage3_max_live_parameters→ 不 OOM 但吞吐差就查stage3_prefetch_bucket_size和stage3_max_reuse_distance。开了 CPU/NVMe 卸载offload_param/offload_optimizerdevice设为cpu或nvme后前两项的下探空间更大。Zero-3 常见报错的排查清单使用参数时断言失败或挂起参数在定义模块之外被访问共享 embedding、别的模块引用 norm 权重DeepSpeed 不知道要提前收集。用register_external_parameter(module, param)登记外部依赖。吞吐明显下降、通信占比高预取桶或重用距离调得过小gather-释放过于频繁。先恢复默认值对比再逐步下调定位。保存出的权重不完整Zero-3 的权重分片散落在各卡上state_dict里拿不到完整张量。把stage3_gather_16bit_weights_on_model_save置为true或在 rank 0 上先 gather 再保存。显存 OOM确认offload_param/offload_optimizer是否配置、device是cpu还是nvme同时检查stage3_param_persistence_threshold是否开得过大导致大量小参数常驻显存。Zero-3 本质是用用到再临时凑齐、用完就释放换取内存空间上面所有配置项都在服务这件事凑得更快、占得更少。从默认值出发先观察显存占用与单步耗时再逐项微调方向就不容易跑偏。【免费下载链接】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),仅供参考