DeepSpeed ZeRO-3 参数分片实战按需聚集机制拆解与 stage3 配置调优指南【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed[rank 1] CUDA out of memory. Tried to allocate 112.00 GiB. Process 0 has 27.97 GiB in use. Of which 26.41 GiB is allocated by PyTorch.微调 70B 模型时,普通数据并行要求每张 GPU 持有完整参数副本,单卡直接 OOM。DeepSpeed ZeRO-3 的参数分区(把参数切成 N 份,每张卡只留自己那 1/N)配合按需聚集,把单卡显存占用从全量副本降到局部分片临时聚集。本文从ZeroParamStatus状态机讲起,覆盖预取、释放全生命周期,以及 7 个 stage3 关键参数的调优边界。训练 1T 参数模型时,每张卡实际驻留的参数只有总量的 1/512,这是 ZeRO-3 与张量/流水线并行的核心差异。官方发布的内存开销对比图如下:为什么 OOM:ZeRO-3 用状态机替代常驻副本数据并行下,参数是常驻副本:每张卡从初始化到结束都持有全部权重。ZeRO-3 反其道而行,把参数状态显式建模为一个三态枚举,每个参数在任意时刻只能处于其中一种状态:# 来源deepspeed/runtime/zero/partition_parameters.py class ZeroParamStatus(Enum): # parameters are fully present and ready for use on all processes AVAILABLE 1 # parameters are either partitioned or remote in some or all process NOT_AVAILABLE 2 # parameters are being gathered. INFLIGHT 3设计决策链在这里很清楚:为什么不做全量加载延迟释放:常驻副本就是数据并行,显存天花板等于模型大小,分区毫无意义。为什么不做隐式懒加载:PyTorch 的Parameter默认随时可读,但 ZeRO-3 下未聚集的参数底层存储只有 0 个元素,误读会得到 shape 为(0,)的张量,前向直接算出 NaN。所以必须用显式状态位(ds_status)在读写路径上拦截。为什么引入 INFLIGHT 中间态:聚集走异步allgather,提交后数据尚未落地。没有第三态就无法防止聚集还没完成就被释放或同一参数被重复提交两类竞态——代码里InflightParamRegistry会对重复注册直接抛RuntimeError(partitioned_param_coordinator.py)。预取、使用、释放:三层钩子怎么挂进前向/反向状态机只是账本,真正驱动状态迁移的是挂在每个子模块上的四组钩子,入口在 parameter_offload.py 的pre/post_sub_module_forward/backward_function。以反向为例:# 来源deepspeed/runtime/zero/parameter_offload.py def pre_sub_module_backward_function(self, sub_module): param_coordinator self.get_param_coordinator() param_coordinator.trace_prologue(sub_module) if param_coordinator.is_record_trace(): param_coordinator.record_module(sub_module) param_coordinator.fetch_sub_module(sub_module, forwardFalse)fetch_sub_module内部做两件事:把本模块NOT_AVAILABLE的参数批量 allgather 提交(forward_fetch_submit与forward_prefetch_submit是两个独立计时项,见 coordinator 类常量),然后从预取队列里弹出下一个即将使用的模块提前发起通信。这就是预取的落点——不是猜,而是第一步前向时录制的模块访问顺序(trace)决定的,ZeRoTraceMode会在录制/复用/失配三种模式间切换。两个值得注意的工程细节:通信 dtype 自动降级:get_allgather_dtype检测参数上是否打了comm_dtype标记(autocast 场景),有则用 bf16/fp16 通信,没有则按原 dtype,通信量和显存峰值都受它影响。释放带流同步保护:free_param在丢弃存储前先record_stream,防止 CUDA 异步执行时内核还在读这块显存:# 来源deepspeed/runtime/zero/partition_parameters.py if get_accelerator().on_accelerator(param.data): if not get_accelerator().is_synchronized_device(): param.data.record_stream(get_accelerator().current_stream()) # param.data doesnt store anything meaningful in partitioned state param.data torch.empty(0, dtypeparam.dtype, deviceparam.device) param.ds_status ZeroParamStatus.NOT_AVAILABLE数据佐证:512 卡 1T 参数,49 TFLOPS/卡ZeRO-3 比数据并行多出约 50% 的通信量(每层参数前后各一次 allgather),官方实现靠计算-通信重叠把它几乎全部藏进计算时间里。公开数据(引自 2021-03-08-zero3-offload.md):对比维度ZeRO-2 OffloadZeRO-3 Offload单 V100(32GB)1.5TB 内存可训模型13B40B(约 3 倍)训练 1T 参数所需 GPU 数不可行512 张 V100(3D 并行需约 1600 张)1T 参数持续算力(512 V100)—25 PFLOPS,约 49 TFLOPS/卡相对 ZeRO-2 Offload 单卡吞吐基线13B 模型上约 2 倍7 个 stage3 配置项:默认值、作用与调错后果所有项定义在 zero/config.py,按配置项逐行给出:配置项默认值作用调错会怎样stage13 才启用参数分区停留在 2 只分区梯度优化器状态,70B 依然放不下stage3_prefetch_bucket_size5e7异步 allgather 预取的桶大小(元素数)太小→通信碎、带宽利用率低;太大→临时聚集显存峰值升高stage3_max_live_parameters1e9允许同时驻留 GPU 的参数元素上限调大省通信、费显存;OOM 时优先调它stage3_max_reuse_distance1e9两次使用距离超过该值就释放参数调小→重算场景反复 gather;调大→显存堆积stage3_param_persistence_threshold1e5小于该规模的参数永久驻留不分区默认值对大量小 bias 已经有效,一般不用动offload_param.devicenone分片卸载到cpu/nvme设为 cpu 后吞吐下降约 20%-30%,换取数倍显存空间stage3_gather_16bit_weights_on_model_savefalse存 checkpoint 时自动聚齐全量权重不开则拿不到可直接加载的完整权重文件动手链路:从单行配置到生产级第 1 步,最简配置——只加一行stage: 3:{ zero_optimization: { stage: 3 } }做错了会怎样:模型超过单卡显存时,deepspeed.initialize之后第一次前向就会 OOM——因为 8B 以上模型不能直接在 GPU 上完整构造,需要配合第 3 步。第 2 步,模型构造期就分区——用zero.Init上下文,参数一分配就切片到各卡:# 来源docs/_tutorials/zero.md with deepspeed.zero.Init(data_parallel_groupmpu.get_data_parallel_group(), remote_deviceget_args().remote_device, enabledget_args().zero_stage 3): model GPT2Model(num_tokentypes0, parallel_outputTrue)做错了会怎样:不用zero.Init直接model.to(device),完整权重先在 GPU 上存在一次再分区,70B 在这一步就 OOM。第 3 步,生产级叠加卸载与调参:{ zero_optimization: { stage: 3, // 参数梯度优化器状态全分区 contiguous_gradients: true, // 梯度连续化,减少碎片 stage3_prefetch_bucket_size: 1e7, // 预取桶,通信带宽不足时调大 stage3_max_live_parameters: 1e9, // 显存紧张时调小 stage3_param_persistence_threshold: 1e5, // 小参数常驻 stage3_gather_16bit_weights_on_model_save: true, offload_optimizer: { device: cpu }, // 优化器状态卸载 offload_param: { device: cpu } // 参数分片卸载,最后手段 } }调参顺序建议:先stage3_max_live_parameters(显存旋钮),再stage3_prefetch_bucket_size(通信旋钮),offload_param放最后——CPU 带宽是 PCIe 的零头,每多卸一层吞吐损失越明显。三个踩坑实录⚠️坑一:模块外读权重,前向输出全 NaN。现象:某个模块在forward里用到了兄弟模块的weight,结果算出 NaN,单步调试发现该参数ds_status NOT_AVAILABLE,读到的数据是 0 元素张量。 根因:ZeRO-3 只协调属主模块内的参数访问,跨模块引用不在它的依赖图里,没人负责为它发起 gather。 修复:在需要该参数的模块上注册外部依赖,register_external_parameter会把它并入该模块的前向/反向 gather 范围:# 来源deepspeed/runtime/zero/partition_parameters.py class ModuleZ3(torch.nn.Module): def __init__(self, *args): super().__init__(*args) self.layer1 SomeLayer() self.layer2 OtherLayer() deepspeed.zero.register_external_parameter(self, self.layer1.weight)顺带一提:如果参数是挂在OrderedDict里、且访问发生在前向内,ZeROOrderedDict.__getitem__会自动替你注册并 gather(parameter_offload.py 第 82 行),这是为什么有些代码没注册也能跑。坑二:Cannot free a ZeRO-3 parameter while it is still active in submodules。现象:训练中途free_param抛RuntimeError。 根因:在GatheredParameters上下文里修改了参数,但modifier_rankNone,各 rank 的副本状态不一致,属主模块还没释放,coordinator 就尝试回收。 修复:给GatheredParameters(..., modifier_rankrank)指定一个 rank 作为修改者,由它广播更新,保证所有副本一致后再释放。坑三:log_trace_cache_warnings持续告警,预取失效。现象:日志反复出现 trace 与当前 forward/backward 不匹配,吞吐量明显下降。 根因:训练循环里用了动态控制流(按 batch 长度走不同分支),第一步录制的模块访问顺序(trace)不再覆盖实际路径,预取退化为同步 gather。 修复:固定控制流路径;或打开log_trace_cache_warnings: true观察失配频率,必要时改用stage3_max_reuse_distance放宽释放策略兜底。行动出口 可直接执行的命令序列:git clone https://gitcode.com/GitHub_Trending/de/DeepSpeed cd DeepSpeed pip install -e . deepspeed --num_gpus8 train.py \ --deepspeed zero3_config.json \ --deepspeed_autotuning_profile # 先用默认 3 跑通,再按上表逐项调进一步阅读路径(均为仓库内相对路径):官方教程与 JSON 示例:docs/_tutorials/zero.md全部 zero 配置项说明:docs/_pages/config-json.md状态机与聚集实现:deepspeed/runtime/zero/partition_parameters.py钩子挂载与外部参数自动注册:deepspeed/runtime/zero/parameter_offload.py发布数据与 3D 并行对比:docs/_posts/2021-03-08-zero3-offload.md【免费下载链接】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），仅供参考