DeepSpeed ZeRO-3：参数怎么分、什么时候加载

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训练日志停在第 47 层共 125 层torch.cuda.OutOfMemoryError: Tried to allocate 2.37 GiB。机器上有 8 张卡单卡显存压力却一分没少——数据并行里每张 GPU 都存着完整模型副本加卡并不省显存。DeepSpeed ZeRO-3 改变的就是这一点参数、梯度、优化器状态全部切分给所有 rank算到哪一层才把那一层的参数临时取过来用完立刻释放。它到底在做什么参数的存放、获取与释放参数放哪里每张卡只存 1/N 的分片传统数据并行下参数、梯度、优化器状态在每张卡上各存一份完整冗余。ZeRO-3 把三者都按 rank 切分deepspeed.initialize之后每个参数被切成 N 份当前 rank 只保留自己那份 1/N其余部分躺在其他卡上。你在本地模块里看到的param.data其实是一个 0 元素的空张量真正的分片存在框架挂上去的ds_tensor里。拿 7.5B 参数模型、Adam 优化器、64 张卡做例子纯数据并行每卡约需 120GBZeRO-3 分区之后降到约 1.9GB——参数、梯度、优化器状态全部被卡数摊薄。代价写在同一张图的右列每次前向/反向都要做 all-gatherZeRO-3 的通信量约为纯数据并行的 1.5 倍。这笔开销能否被藏住取决于下面两个机制。参数何时可用算到前才临时 gather参数生命周期由一个三态状态机控制class ZeroParamStatus(Enum): AVAILABLE 1 # 完整参数已在本地 GPU可直接计算 NOT_AVAILABLE 2 # 只有分片本地无完整副本 INFLIGHT 3 # gather 正在进行中模块前向开始前框架检查该模块的参数凡是NOT_AVAILABLE的就发起异步 allgather状态置为INFLIGHT计算流轮到它使用时同步一下再翻回AVAILABLE。gather 走通信流async_opTrue第 k 层的参数往往在第 k-1 层算完前就到位了通信被计算掩盖。16 位精度下框架还会根据精度配置自动选择通信 dtype传输体积不会大于模型本身的存储精度。整套逻辑在 deepspeed/runtime/zero/partition_parameters.py预取窗口和 gather 的数据类型选择都能在这个文件里顺藤摸到。参数什么时候释放用完立刻还给分片一层的前向/反向结束后该层参数马上被释放def free_param(param): param.data torch.empty(0, dtypeparam.dtype, deviceparam.device) param.ds_status ZeroParamStatus.NOT_AVAILABLE即param.data缩回空张量、显存归还状态退回NOT_AVAILABLE下次要用再取。这里有个硬约束如果参数还处在某个子模块的激活使用状态里计算还在引用它free_param会直接抛RuntimeError而不是悄悄释放。这个设计避免了参数还在用就被释放的最坏情况但同时也带来实践中最常见的一个报错后面按现象排查时会讲。配置不是抄模板四个参数的联动关系ZeRO-3 里最常调的四个 key 彼此不是独立的stage3_max_live_parameters能囤多少已 gather 的参数是显存硬上限stage3_max_reuse_distance预取的前瞻窗口预计某参数还有多远会被复用就提前多远去取stage3_prefetch_bucket_size单次预取批的大小stage3_param_persistence_threshold小于该阈值的参数常驻 GPU完全不参与 gather/释放循环。前两个决定热内存囤多少调大省通信但逼近 OOM调小省显存但 gather 循环更频繁、通信更忙。持久化阈值最容易被忽略——Transformer 里 qkv、LayerNorm、bias 这类小参数数量极多让它们走一遍 gather 流程纯属浪费阈值设小了白设设大了又会挤占 live 参数的空间。完整配置项可对照 docs/_tutorials/zero.md。起步配置百亿级模型、8 卡、纯 GPU{ zero_optimization: { stage: 3, contiguous_gradients: true, stage3_prefetch_bucket_size: 5e7, stage3_param_persistence_threshold: 1e5, stage3_max_live_parameters: 1e9, stage3_max_reuse_distance: 1e9 } }激进配置大模型 CPU offload{ zero_optimization: { stage: 3, stage3_prefetch_bucket_size: 2e7, stage3_param_persistence_threshold: 1e4, stage3_max_live_parameters: 1e7, offload_param: { device: cpu, pin_memory: true }, offload_optimizer: { device: cpu, pin_memory: true } } }两组值的差异全在 offload 路径上参数落盘到 host 后单次 gather 的带宽比卡间通信慢一个量级大桶装不下也藏不住延迟只能把桶和常驻参数一起缩小用空间换时间pin_memory则是给 H2D 拷贝做的。offload 逻辑实现在 deepspeed/runtime/zero/parameter_offload.py。跑起来之后的事按现象排查报错 Cannot free a ZeRO-3 parameter while it is still active in submodules这就是上面释放守卫被触发的结果同一个参数被两个模块引用典型是共享 embedding、输入输出 tied weight第一个模块算完尝试释放第二个还在用。修法看参数的用法。如果是某模块引用了外部模块的参数在构建期把依赖显式注册给框架让它计入释放判断class Decoder(torch.nn.Module): def __init__(self, shared_embedding): super().__init__() self.emb shared_embedding deepspeed.zero.register_external_parameter(self, shared_embedding.weight)如果只是在代码某处临时访问分片参数比如保存时用deepspeed.zero.GatheredParameters(param, modifier_rankrank)包住使用段即可注意修改参数时必须指定 rank 做广播否则各 rank 状态会不一致。通信占比高比 ZeRO-2 明显慢预期内不是 bug。ZeRO-3 通信量约为数据并行的 1.5 倍小模型、小 batch、高频 gather 场景下差距最明显。调整顺序先拉持久化阈值让常驻小参数留下来再看预取窗口是否太小导致 gather 迟到最后查网络拓扑有没有瓶颈。值得知道的是模型越大这笔开销占比越低仓库官方博客docs/_posts/2021-03-08-zero3-offload.md测得 512 张 V100 上400B、1T、2T 参数模型的单卡吞吐分别为 50.2、48.9、40.7 TFlops到 2T 只掉约 19%——计算时间足够长gather 能被完全掩盖。反过来模型小、单步计算短通信暴露比例自然高这时候应该认真考虑留在 ZeRO-2。检查点里只有分片拿不到完整权重ZeRO-3 的 checkpoint 本来就是按 rank 存的分片属正常现象。要完整 16 位模型配置里加stage3_gather_16bit_weights_on_model_save: true保存时调用model_engine.save_16bit_model(...)框架会把所有分片聚合成完整参数再落盘。已有的 checkpoint 也可以用deepspeed/utils/zero_to_fp32.py离线合并。边界与进阶ZeRO-3 不是万能药。模型在 10B 参数以下、单卡显存又不吃紧时gather 开销可能让你比 ZeRO-2 更慢CPU offload 时 PCIe 带宽是天花板模型不够大时省下的显存换不回拷贝的开销动态控制流复杂、或手动切分参数喂给计算的代码会破坏自动 gather/释放的假设得用GatheredParameters手动接管参数生命周期。再往深走有三个方向值得看ZeRO-3 叠加张量并行Megatron 风格或 auto-tp数据并行 张量并行二维组合面向万亿参数host 内存也不够时继续 offload 到 NVMeZeRO-Infinity 路线DeepSpeed 的 DeepNVME 组件为此而生量化 gather 之类的通信压缩ZeRO 路线直接降低 all-gather 的传输体积拿一个小模型把配置改掉完整跑一遍、看一眼显存曲线比读十篇文章都管用。【免费下载链接】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），仅供参考