DeepSpeed ZeRO实战：从显存瓶颈到大模型高效训练

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 字节/参数是最理想的静态模型状态，实际工程中还会出现：

如果不用混合精度，全部采用 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/s，CPU 内存带宽几十 GB/s，NVMe 则更低。所以 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 时，这些数据是定位问题的第一手线索。很多人加载失败后连自己这次运行的显存基线都没有，只能靠猜，效率极低。先记录，再调参，比什么参数推荐都管用。