★ FEATURED ARTICLE 1. 从显存崩溃说起DeepSpeed当年要解决的真实困境如果你是从大模型爆发那两年才入行的大概很难想象几年前为跑一个 7B 参数模型要经历什么样的折腾。模型本身不大但训练它需要的显存却大得离谱算力倒是其次显存才是真正的墙。我第一次在单卡上尝试全参数训练一个中等规模语言模型时结果非常直接——OOM连梯度累积的机会都不给。这就是 DeepSpeed 当初最想解决的痛点。它不是一个用来炫技的框架而是把“如何在有限显存里塞进更大模型”这件事做到了机制层面的极致。回看这几年它的演进核心其实一直没有变用分片思维消除冗余用分层级的技术手段去换显存、省带宽、提吞吐。这篇内容适合几类人看想跑大模型但显存不够的工程师在做分布式训练调优的研究生以及那些听了很多次 ZeRO 但一直没搞清楚它到底“分”了什么东西的开发者。我尽量把原理和实战都讲透不绕弯子。1.1 大模型训练的真正瓶颈不是算力是显存一张训练卡最短缺的资源往往不是算力峰值而是显存容量。模型参数要放在显存里梯度要放在显存里优化器的状态也要放在显存里但很多人只盯着参数量算账忽略了一个关键事实训练状态远比模型参数本身“重”得多。拿 7B 模型举例。如果用混合精度训练参数通常以 FP16 存储每个参数占 2 字节梯度也常保留 FP16 副本再占 2 字节而 Adam 优化器要保存 FP32 的 master weights、一阶动量、二阶动量每个方向都是 4 字节加起来 12 字节。仅仅这些“模型状态”分摊到每个参数就是 16 字节7B 参数对应的显存需求已经超过 112GB。112GB 是什么概念即便放到现在单张 A100 80G 也放不下更别提当时主流的 32G、40G 显卡。而且这还没算激活值、临时缓冲区、通信缓冲这些运行时的额外开销。真正训练起来同样规模的模型轻松吃掉两三百 GB 显存。所以你会看到很多看起来很矛盾的现象一个能把推理跑得很流畅的模型训练时却动不动爆显存一批小批量大小的训练脚本在单卡上跑不动但推理脚本却轻巧得很。原因就在于训练过程多了一整套优化器状态和梯度生命周期。1.2 传统分布式方案哪里不够用面对“一张卡放不下”的问题传统的解决办法无非是数据并行、模型并行、流水线并行。数据并行最简单每张卡都复制一份完整模型和优化器状态各自处理不同的数据切片然后在梯度计算完后做一次全归约同步。问题是它的显存开销是随卡数线性叠加的——8 张卡就有 8 份完整的模型状态。训练 7B 模型单卡需要 112GB数据并行在 8 张卡上就需要 8×112GB 的总显存而模型本身的容量需求并没有因为卡多了而下降只是你能同时处理更大的 batch 而已。模型并行和张量并行把模型切到不同卡上确实能降低单卡压力但它要求切分后的每一层算子之间频繁通信通信量巨大而且对层设计有约束工程复杂度高。流水线并行可以把不同层放到不同卡上但会出现只让一张卡在某个时刻满负荷的情况利用率和负载均衡一直是痛点。这些方案的本质是“从计算角度切分模型”而 DeepSpeed 的 ZeRO 系列走的是另一条路它不把计算切得粉碎而是让每张卡都只保存一份“数据并行全局状态”的分片。换句话说数据并行的逻辑语义不变但每张卡不再冗余持有完整状态。这个思路上的转变才是后续一系列能力演进的地基。2. ZeRO机制拆解为什么“分片”能换来百倍收益ZeRO 在 DeepSpeed 里不是一个单一功能而是一整套递进式优化策略。很多人听过“ZeRO Stage 1/2/3”的说法但如果你只知道这是“分片优化器状态”那你还是低估了它的设计。我拆开讲一遍每个阶段到底做了什么以及为什么“分片”这个动作能带来这么大的收益。2.1 训练状态到底吃掉了多少显存在深入阶段细节前先把账算清。前面提到的 16 字节/参数是最理想的静态模型状态实际工程中还会出现状态类型存储精度每参数占用说明模型参数副本FP162 字节推理和正向计算使用梯度副本FP162 字节反向传播后用于更新也可能以 FP32 形态存在优化器主权重FP324 字节Adam 需要一份稳定精度的权重副本一阶动量 mFP324 字节Adam 的滚动平均二阶动量 vFP324 字节Adam 的自适应学习率分母如果不用混合精度全部采用 FP32 训练那模型参数和优化器状态会更膨胀每参数至少 20 字节起步。这也是为什么混合精度在现代训练中几乎是标配它不只是省一半参数显存更重要的是配合分片策略后优化器状态成为显存账本上的“大头”而 ZeRO 的阶段性设计正好先从这块下手。除了静态状态训练时的激活值也不可忽视。一个深层的 transformer 模型正向过程会保存大量中间激活用于反向传播这部分显存随序列长度和 batch size 近似线性增长。这也是后来 gradient checkpointing、Activation Offload 等机制被整合进 DeepSpeed 生态的原因。2.2 三阶段递进从优化器分片到全分片Stage 1只分片优化器状态。这是 ZeRO 的起点。优化器状态在模型状态中占比最大Adam 的三份 FP32 状态占 12 字节/参数占静态状态的 75%。Stage 1 把优化器状态按 rank 切片每个进程只保存一份全局优化的分片更新时每个秩只更新自己那部分再通过通信确保参数变化一致。Stage 2分片优化器状态和梯度。梯度在反向传播完成后、更新之前会被规约成全局梯度这个全局梯度在数据并行下每张卡都保留一份完全是冗余的。Stage 2 把梯度也切片每个秩只保留对应自己优化器分片的那一部分梯度。这样显存又降一截。Stage 3参数也分片。到了这一层每张卡上连完整模型参数都不保留了只保存参数分区。每次需要前向或反向计算某个层时该层的权重需要先通过 all-gather 从各个分区汇聚到对应计算设备上算完再丢弃。这几乎消除了数据并行下的所有静态冗余。用一个不太严谨但直观的比喻数据并行模式下相当于 8 个人各背了一套完整的露营装备Stage 2 相当于每个人只背自己那份帐篷杆、防潮垫和睡袋但整套装备合起来仍然完整Stage 3 则更进一步连每件工具都不重复持有需要哪个工具就去对应的人那里借用完还回去。这里有个很重要的代价认知Stage 3 并非免费每层参数的 all-gather 会增加大量通信但它的收益是显存随卡数近似线性下降理论上可以让单卡容纳百亿甚至千亿参数规模。这也是为什么它在超大模型训练场景中如此重要。2.3 通信开销与Offload的取舍ZeRO 的每一步优化都在“显存容量”和“通信开销”之间做交换。从实测经验看Stage 3 的通信量大约是普通数据并行的 1.5 倍左右远低于大规模张量并行那种“每个算子都要卡间同步”的通信压力所以它在多机场景下反而更容易扩展。Offload 则是另一条交换路径把原本要放在显存里的状态转移到 CPU 内存甚至 NVMe 磁盘上。CPU offload 适合优化器状态和参数副本NVMe offload 则能在极端情况下把模型状态放到磁盘上。代价是带宽差异显存带宽通常在几百 GB/sCPU 内存带宽几十 GB/sNVMe 则更低。所以 Offload 适合“只要能跑通、不追求极致速度”的场景比如消费级显卡在中大规模模型上的微调。我的建议是如果训练吞吐是第一目标优先用纯 Stage 2别轻易开 Offload如果目的是验证大模型算法可行性Stage 3 CPU Offload 是性价比最高的组合。很多人一上来就开 Stage 3 Offload 反而把速度拖到不可接受其实往往不是框架问题而是没搞清楚取舍方向。这一点在后面实战部分还会细说。3. 能力的边界扩展从纯训练演进到推理优化DeepSpeed 被大众熟知是因为 ZeRO 在训练领域的成功但它的演进远不止训练。回看整个技术路径DeepSpeed 在推理端也做了大量工作而且这些工作对今天的部署实践影响深远。如果只把它当“训练加速库”你会错过它半张能力版图。3.1 Inference引擎做了什么训练好的大模型部署到线上时最直接的问题是吞吐上不去。推理阶段没有梯度也不需要优化器状态显存压力比训练小很多但计算模式完全不同它要求极低延迟、高吞吐并且要处理动态 batch、变长输入等现实情况。DeepSpeed Inference 的思路不是简单封装几个算子而是把 transformer 层里适合融合的 kernel 自动做算子融合减少 GPU kernel 启动开销同时对权重进行自动切片让多个 GPU 协同推理。一般的卷积或者矩阵乘小算子每次调用都有固定开销融合后数据在寄存器与共享内存之间流动更高效实测吞吐提升往往以倍数计。更关键的是它把训练阶段的“分片思维”延伸到了推理。这就是所谓的 ZeRO-Inference加载大模型时不同设备只持有模型参数的分片前向计算时按需聚合。这样单机多卡可以跑千亿级模型推理不需要像传统张量并行那样对模型结构做精细切分部署复杂度低了一大截。3.2 与训练框架生态的整合演进DeepSpeed 的长期演进中还有一个明显趋势不再只做“显存优化”而是逐步成为一个覆盖训练全流程的系统。从混合精度管理、学习率调度、梯度累积到自动调参、断点检查、模型并行策略选择它都想给出一套工程化答案。在注意力计算方面它对前几年流行的 Flashattention 类优化做了深度适配自动选择最合适的注意力内核让长序列训练吞吐再上一个台阶。在稀疏化方向上它支持混合专家模型MoE的训练与推理不同 token 可以通过门控网络只激活部分专家配合“专家并行”让稀疏模型在训练成本上比稠密模型低得多。此外整套框架也把 RLHF 场景纳入能力图谱支持在训练过程中串联多阶段任务。这类能力让同一套代码基础设施能从头训练、微调到对齐训练复用不用反复切换工具。本质上DeepSpeed 的演进路径是从“解决显存问题”走向了“解决整个大模型生命周期问题”。4. 一份可上手的ZeRO-3训练配置实操讲了这么多原理接下来给一份可以直接落地的训练配置。我以单机多卡场景为例模型规模中等偏大显存不够但追求尽可能高的吞吐。下面这份 ZeRO-3 配置是我在多轮实验中最终稳定下来的版本参数都不是拍脑袋写的。4.1 环境准备与初始策略选择DeepSpeed 的安装不复杂但有两个地方容易踩坑一是需要和你的训练框架版本对齐二是有些扩展算子需要编译编译器环境不全会导致运行时才报错。建议先跑一遍环境自检命令确认当前环境状态。然后启动训练的方式与普通脚本略有不同不再直接调用 python而是通过 deepspeed 启动器分发任务训练脚本内需要额外传入配置文件路径。pip install deepspeed # 检查环境确认 CUDA、编译器、通信库等核心组件状态 ds_report # 单机多卡或单卡启动示例 deepspeed train.py \ --deepspeed \ --deepspeed_config ds_config.json初始策略选择上我的经验很简单如果你的模型在单卡上能勉强塞进静态权重并且训练 batch 不需要很大优先用 ZeRO Stage 2它的速度和稳定性最好如果权重本身已经超过单卡显存直接上 Stage 3但要做好通信调优如果目标是把消费级显卡利用起来考虑 Stage 3 CPU Offload。4.2 关键配置项逐个拆解下面是一份可运行的 ZeRO-3 配置文件我标注了每个核心字段的含义{ train_batch_size: 1024, train_micro_batch_size_per_gpu: 8, gradient_accumulation_steps: 16, optimizer: { type: AdamW, params: { lr: 1e-4, betas: [0.9, 0.999], eps: 1e-8, weight_decay: 0.01 } }, zero_optimization: { stage: 3, offload_optimizer: { device: cpu }, offload_param: { device: cpu }, overlap_comm: true, contiguous_gradients: true, reduce_bucket_size: 5e7, stage3_prefetch_bucket_size: 5e7, stage3_param_persistence_threshold: 1e5 } }这里有几组概念需要重点理解batch 三件套。train_batch_size 是全局有效 batch它等于 train_micro_batch_size_per_gpu 乘上 GPU 数量再乘上梯度累积步数。三者关系是可互换的例如想要更大全局 batch 而单卡显存有限就调高 gradient_accumulation_steps。offload 选择。我同时打开了优化器和参数的 CPU Offload这是为了在中等显存卡上跑更大的模型。但要注意如果卡显存还能撑住优先只打开 offload_optimizer参数 offload 会让前向计算频繁从 CPU 拉权重吞吐损失明显。通信重叠。overlap_comm 让参数聚合和计算重叠能有效掩盖一部分通信延迟contiguous_gradients 把零散的梯度张量规整成连续缓冲区减少通信碎片化。这两个开关通常建议直接打开。bucket 与阈值。reduce_bucket_size 控制梯度规约的粒度值越小通信越碎片化stage3_prefetch_bucket_size 控制预取权重时每次拉取的参数块大小stage3_param_persistence_threshold 决定多大的参数留在本卡缓存中不释放。这几个值需要按模型结构调整我给的偏保守数值在大多数场景都能稳定运行。4.3 训练过程中的显存与权重处理配置好只是第一步真正训练时还要注意两个工程细节。第一配合 gradient checkpointing 使用。ZeRO 分片解决的是静态模型状态而激活值依然是显存小杀手。开启梯度检查点可以把正向激活按需重算显著降低峰值显存。但代价是约 30% 的计算重开销。对显存紧张的项目来说这是比裁剪 batch 更划算的取舍。第二权重保存与加载要格外小心。ZeRO-3 下默认的 checkpoint 是按分片保存的每张卡只落一个自己的状态文件直接加载到别的策略下会报形状不匹配。官方提供了合并权重的脚本比如将分片权重转化为完整权重后再加载。这个步骤很多人第一次用都会漏掉导致“为什么明明能看到 checkpoint 却加载失败”的困惑。我一般习惯在每轮训练结束或者关键节点保存两份一份是原始分片 checkpoint 用于断点续训一份是合并后的完整权重用于下游任务验证。虽然占磁盘空间但能避免很多返工。5. 实战踩坑记录与调优体会踩坑才是真正区分“读过文档”和“跑过训练”的地方。我把自己在 DeepSpeed 实战中遇到的几类高频问题整理出来这些问题在官方文档里都有简短说明但真正定位时比想象中曲折得多。5.1 三类高频坑的完整排查链路第一类启动阶段就 OOM。现象是训练进程刚开始加载模型就报显存不足连日志都没跑出来。很多人第一反应是减小模型或 batch但真正原因往往是 stage3 的参数持久化阈值设置太大了。该参数默认是较大的值会让大量小参数留在本卡显存而不释放如果模型包含大量 embedding 等参数表就会在启动瞬间耗尽显存。排查顺序先把 stage3_param_persistence_threshold 调小比如降到 1e4 或 1e5再检查是否同时在多卡加载了冗余的优化器状态最后看启动时是否每个 rank 都执行了全量参数初始化。通常第一步就能解决。第二类开 Offload 后速度断崖下跌。现象是开启 CPU Offload 后每个 step 耗时暴涨从秒级变分钟级。问题根源通常是 reduce_bucket_size 或 prefetch bucket 太小导致参数和梯度在 CPU 与 GPU 之间碎了太多次每次都产生传输延迟。排查顺序先把 bucket 调大到 5e7 以上然后确认 CPU 端是否有足够内存带宽双路服务器和普通台式机差异很大最后尝试关闭 offload_param只保留 offload_optimizer对比速度变化。实测中很多人关掉参数 offload 后速度回来了而显存占用并没增加太多。第三类换了 Stage 后加载旧 checkpoint 失败。现象是训练到一半想从 ZeRO-2 切到 ZeRO-3或者反过来结果权重加载报形状错误。原因很直接分片策略变化导致状态字典结构不一致。这不是资源问题是保存格式问题。排查顺序如果当时保存的是分片 checkpoint先手动用脚本合成完整权重确认合成后的 key 名称与模型定义一致如果模型结构里有命名不一致的参数需要先做 key mapping 再加载。防止这类问题的根本办法是预设固定的 stage 并尽量不切换实在要切换就重新加载完整权重冷启动。5.2 我的调优顺序与方法很多资料会直接给一套“默认配置”但实际训练中最好按下面的顺序逐步放宽约束而不是一开始就追求极限第一步先用尽量保守的配置把任务跑通例如 batch 调小、关掉 offload、用 Stage 2。先把正确性验证完再谈效率。第二步逐步提高全局 batch同时观察单卡显存曲线确认没有隐性的显存尖峰。第三步确认显存有余量后再考虑打开 stage3 或 offload但必须单独验证性能回退幅度记录 step 耗时和通信占比。第四步针对通信热点调 bucket、开 overlap_comm并用监控工具观察 GPU 利用率是否有明显空窗口。这套顺序的核心原则是先保证“能跑”再优化“跑得快”最后再做“跑得更大”。很多人直接抄一份 Stage3 Offload 配置跑到最后发现速度灾难性下降回头又很难定位是因为哪一项配置引起的就是因为跳过了中间验证步骤。我自己的经验是DeepSpeed 在大多数情况下的默认值偏向“稳妥”而不是“最优”。比如 reduce_bucket_size 如果模型很小默认值会让通信微型化如果模型很大默认值又可能让通信粒度太粗导致某张卡等别人的梯度等了很久。这类参数必须结合自己的模型计算图去测没有一劳永逸的万能值。最后再分享一个小技巧训练脚本里最好把每次启动的实际显存峰值、step耗时、通信占比全部打进日志。遇到偶发卡顿或 OOM 时这些数据是定位问题的第一手线索。很多人加载失败后连自己这次运行的显存基线都没有只能靠猜效率极低。先记录再调参比什么参数推荐都管用。
阅读完成 · 觉得有帮助？
