1. 显存墙面前分布式训练被逼出来的十年十年前当深度学习圈还在用单卡跑VGG和ResNet的时候大概没人会想到十年后的今天一个 1750 亿参数的大模型会让学术界和工业界集体患上显存焦虑症。那时候我们讨论的是怎么把 Batch Size 调大一点现在讨论的是怎么把 80GB 的 H100 拼成一台能装下十万亿参数的虚拟机器。这十年里硬件算力按照摩尔定律的余晖缓慢爬升但模型参数量的增长速度却是指数级的——这两条曲线之间的剪刀差就是分布式训练框架存在的全部理由。在这个赛道上DeepSpeed 是我个人认为最值得拿出来做十年复盘的项目之一。它从 2020 年开源时的显存优化工具逐步演变成一个覆盖训练、推理、压缩、服务化的全栈系统。它不是唯一的选择PyTorch 原生的 DDP/FSDP 也在不断进化Megatron-LM 则在张量并行和流水线并行上做得极其深入。但 DeepSpeed 最特别的地方在于它的设计哲学不要求用户重写模型代码尽量在框架层把显存和通信的潜力榨干。这篇文章我不会泛泛地吹嘘某个框架多好用而是想以一个踩过无数坑的实践者视角把 DeepSpeed 这十年的演进脉络拆开——它到底解决了什么问题每个核心技术的原理是什么实际用起来有哪些文档里没写的坑以及它和 PyTorch 生态之间的复杂关系。先放一张我根据自己的使用经验总结的粗略时间线方便对这段历史不熟的读者建立坐标系时间阶段核心痛点DeepSpeed 的代表动作2019-2020单卡显存不够模型并行无从下手开源 ZeRO 系列打出口号训练 1000 亿参数模型2020-2021显存省下来了但通信耗时训练吞吐上不去ZeRO-Offload、ZeRO-Infinity把 CPU 内存和 NVMe 都变成显存2021-2022模型结构越来越复杂训练和推理割裂支持稀疏注意力、混合专家模型收敛到 Transformer 架构2022-2023大模型推理成本高企DeepSpeed-Inference、MII把推理优化纳入体系2023-至今系统复杂度爆炸普通团队跑不动大模型DeepSpeed-Chat、Autotuning、PerfZero自动化配置和基准化如果你只打算把这篇文章当作一份 DeepSpeed 的论文笔记那可能有点亏。我更想传达的是这十年的演进路径本质上是一条把分布式训练从天才的玩具变成普通工程师的日常工具的路线图。接下来我会把这条路线上的关键站台一站一站地展开。2. 为什么偏偏是显存成了墙DDP 时代的痛与 ZeRO 的破局要理解 DeepSpeed 为什么一出来就引起轰动得先回到它出现之前那个黑暗时代看看。2020 年前后大家训练大模型的标准做法是什么PyTorch 的 DistributedDataParallelDDP已经相当成熟了它的逻辑很简单每个 GPU 上都放一份完整的模型副本前向传播各自算反向传播时用 AllReduce 把所有 GPU 上的梯度求和然后各卡拿到的都是平均梯度各自更新参数。这套逻辑在模型 10 亿参数以内的时候是够用的。一张 32GB 的 V100 差不多能塞下一个 10 亿参数的模型加上优化器状态和激活值。但一旦模型规模跨过百亿、千亿参数DDP 就彻底失效了。原因不是通信带宽不够而是你压根没办法把模型塞进任何一张卡里。我们可以简单算一笔账。一个 10 亿参数的模型以 FP16 存储光模型权重就需要 2GB。Adam 优化器需要维护两份状态——动量momentum和方差variance这两份在混合精度训练里通常用 FP32 保存也就是每份 4GB两份就是 8GB。再加上梯度FP16 的 2GB但 ZeRO 出现之前很多实现里梯度也按 FP32 算那就是 4GB还有前向过程中保存的激活值、临时缓冲区、通信缓冲区……你可以看到模型参数只占了一切显存需求的一小部分绝大多数显存其实是被优化器状态和梯度吃掉的。DDP 的做法是每个 GPU 都重复保存这些状态显存需求随着模型大小线性膨胀。模型并行Tensor Parallelism / Pipeline Parallelism虽然是另一条路但它需要你重写模型的层内结构或者手动切分层工程复杂度极高。这时候 ZeROZero Redundancy Optimizer出现它的核心洞察极其简洁既然每个 GPU 上都存了完全一样的优化器状态和梯度副本而这些副本在每次更新迭代中做的事情完全一样那为什么不把它们切碎了存在不同 GPU 上用的时候再汇总呢ZeRO 分了三个阶段来消除这三种冗余Stage 1把优化器状态Adam 的 momentum 和 variance分片到各 GPU。理论上 10 亿参数的模型优化器状态从 8GB 降到 8GB / NN 为 GPU 数。Stage 2把梯度也分片每个 GPU 只保存自己负责的那一段梯度。前向反向结束后用 ReduceScatter 把各卡的梯度归并到对应分片。Stage 3连模型权重都分片前向传播时通过 AllGather 临时收集所需的权重切片反正反向时再收集一次。这就是能被 DeepSpeed 吹上天的模型并行对用户透明的原因。从 DDP 到 ZeRO Stage 2显存占用几乎变成零冗余。我在实际测试里把 30 亿参数的 GPT 模型塞进 4 张 V100每张 32GB跑通了而如果用 DDP这个模型大概需要 12 张才能放下。ZeRO 的本质是把分布式系统中的每卡全量冗余改成全局单副本 通信换显存用通信开销换取显存上限的扩展。3. ZeRO-Offload 与 ZeRO-Infinity把内存和硬盘都变成显存如果故事只到 ZeRO Stage 3 就结束DeepSpeed 可能也就和 FSDP 打个平手谈不上十年演进。真正让我觉得这个团队在做超前设计的地方是他们很快意识到就算你把显存里的冗余全消除了单卡 32GB、40GB 的物理显存上限依然横在面前。当一个团队手里只有几张卡、每张卡只有 32GB 显存却想去跑一个 1750 亿参数的模型时物理显存是无论如何都不够的。于是 2021 年出现了 ZeRO-Offload。这个设计的思路简单说就是不常用的东西搬到更便宜的地方去。哪个地方比 GPU 显存更便宜CPU 内存。哪个地方比 CPU 内存更便宜NVMe SSD。具体实现上ZeRO-Offload 把优化器状态和梯度放到了 CPU 内存GPU 上只保留模型参数和激活值。每一轮迭代计算在 GPU 上完成反向传播算出梯度后把梯度拷贝到 CPU由 CPU 侧的优化器更新参数然后把更新后的参数拷贝回 GPU。听起来很笨拙但它在带宽允许的范围内确实有效。我的经验是当 GPU:CPU 的 PCIe 带宽足够比如 PCIe 4.0 x16 的 32GB/s模型在 10 亿到 50 亿参数这个区段时ZeRO-Offload 的吞吐损失大概是 20% 到 40%换来的却是8 张 32GB 显卡就能跑 100 亿参数模型这种以前想都不敢想的事情。ZeRO-Infinity 则更进一步把模型权重也放到了 NVMe 上GPU 显存只是在计算瞬间才缓存一部分权重切片。这个操作真的把显存墙这个概念本身瓦解了——只要你的硬盘空间足够大理论上可以训练任意规模的模型。我见过有人用 16 张 80GB A100 在 ZeRO-Infinity 下跑万亿参数稀疏模型虽然没有跑满但至少是真实可跑的。当然NVMe 的带宽比 CPU 主存又低一个数量级吞吐损失就相当惨烈了。在纯 NVMe offload 场景下训练速度可能比纯 GPU 显存跑慢 5 到 10 倍。这不是残疾玩法是极端场景下的勇气玩法。这两项技术的定位其实非常清楚Offload 不是为了让你跑得更快而是为了让没有那块顶级硬件的团队也能跑起来。很多刚接触 DeepSpeed 的读者会盯着 benchmark 抱怨Offload 之后速度掉了好多我建议换个角度看待——它解决的是完全跑不了到能跑但慢一点的问题在这个前提下任何吞吐都是赚的。4. 从优化器到全流程引擎自动并行、稀疏注意力与混合专家模型DeepSpeed 早期被贴的标签是ZeRO 优化器但到了 2021 年下半年它的定位开始明显扩大。原因很现实以 Transformer 为骨架的大模型不只是一个内存问题它还有计算模式多样化的挑战。先看自动并行Automatic Parallelism。传统上Tensor Parallelism张量并行和 Pipeline Parallelism流水线并行都需要用户手动切分模型。张量并行要把每个 Transformer 层的 Q、K、V 矩阵按列切开、把输出矩阵按行切开还得在每个注意力前后插入 AllReduce 通信流水线并行则要把整个模型的层切成若干段分配到不同的设备上还要仔细编排 micro-batch 的调度顺序。这些工作异常繁琐且极易出错——哪怕矩阵切分维度写错一个数字都能让训练在第一步就崩掉。DeepSpeed 的自动并行模块想做的事是让用户写出一个普通的模型定义框架自动帮你分析每层的参数量、计算量和通信量决定哪些层该做张量并行、哪些层该做流水线并行。它背后的核心是一个代价模型cost model框架会对每种并行策略做一次前向传播的显存和耗时预估然后选一个综合最优的方案。说实话这个自动并行功能在我实际用下来的体验是方向正确但远未完美复杂的模型结构仍然需要手动干预。但它代表了一个重要的趋势信号——系统的复杂度正在从用户侧转移到框架侧这为后来大模型平民化打下了基础。再看混合专家模型Mixture of ExpertsMoE。MoE 本质上是一种极其浪费的稀疏化方法你有成百上千个 expert 子网络但每次只激活其中极少几个。举个例子一个 MoE 模型有 128 个 expert每个 token 只路由到 top-2那么计算量理论上只有全部激活的 1/64但显存却要装下所有 expert 的权重。这种结构天生就适合 ZeRO Stage 3——不用的 expert 权重分散存在不同设备上用的时候才收集。DeepSpeed 对 MoE 的支持在 2021 年底到 2022 年间逐渐成熟提供了 MoE 层的通信优化比如 expert parallel 的数据路由和负载均衡损失防止所有 token 都堆在同一个 expert 上。我用过的体验是一个 1 万亿参数的稀疏 MoE 模型实际激活参数量可能只有 130 亿用 8 张 A100 就能训练而同样参数量稠密模型是极其不现实的。这种以大存小的模型形态让显存墙问题的重心从能不能装下转移到了怎么高效路由令牌上。5. 并行策略的组合博弈DP、TP、PP、EP 之间到底怎么选很多刚接触 DeepSpeed 的人都会有一个困惑ZeRO、张量并行、流水线并行、专家并行这些听上去都是并行到底有什么区别什么时候应该用哪个这其实是个特别值得认真讲清楚的问题我在实际项目中见过太多人混用概念最后把配置文件改得一团糟。这样理解可能更清晰一点并行策略切分维度通信模式最擅长的场景主要代价数据并行 (DP/DDP)数据样本AllReduce模型能塞进单卡时显存冗余高ZeRO 分片模型状态参数/梯度/优化器分片 AllGather / ReduceScatter模型装不进单卡但希望保持数据并行语义通信量大张量并行 (TP)矩阵行/列AllReduce每层多次单层计算量特别大通信频繁延迟敏感流水线并行 (PP)模型层点对点发送模型层数很深流水线气泡bubble浪费专家并行 (EP)expert 子网络All-to-AllMoE 模型通信模式不规则路由开销实际操作中大部分百亿到千亿参数的稠密模型训练采用的是DP ZeRO Stage 1/2 PP的组合用数据并行撑满通信吞吐用 ZeRO 消除优化器状态的冗余用流水线并行把模型拆到不同设备组上。只有在单卡的算力无法容纳单层需要全量计算的超大矩阵时才会把 TP 加进来——TP 的通信频率是所有并行策略里最高的它对 GPU 之间的互联带宽要求极其苛刻。几周前我在一个项目里尝试训练一个 130 亿参数的稠密模型硬件是 32 张 A100 80GB节点间用的是 InfiniBand。最开始的配置是纯 ZeRO Stage 3吞吐大概 40 TFLOPs per GPU后来把层数切成两组做流水线并行PP2吞吐掉到了 35 TFLOPs。原因很典型流水线并行引入的 bubble空闲等待吃掉了收益而这个模型单层复杂度并没有大到需要跨设备做张量切分的程度。后来我们干脆改回 ZeRO Stage 2 DP吞吐反而回升到 38 TFLOPs。这个例子想说明的是并行策略没有绝对的好坏一切取决于你的模型结构、硬件拓扑和通信库的实测表现。DeepSpeed 的社区里有很多人喜欢无脑开 ZeRO Stage 3但事实上在 2D 并行DP PP就能装下的场景里Stage 3 的通信开销反而可能拖慢训练速度。配置 DeepSpeed 的第一步永远是跑 benchmark而不是抄一台已经调好的配置哪怕它来自官方文档。6. 训练和推理的边界消融DeepSpeed-Inference、MII 与 CPU 卸载的第二春到了 2022 年DeepSpeed 团队做了一个让我印象很深的转向他们不再只啃训练这块硬骨头而是开始把视角扩展到了推理侧。原因也不难理解——训练一个千亿模型可能只会跑几个月但推理部署会用几年推理成本的总量往往比训练更高。DeepSpeed-Inference 的核心动作有两块。第一块是针对 Transformer 结构的 kernel 融合把 Multi-Head Attention 中的多个小算子融合成一个大 kernel减少 kernel launch 的开销和中间张量的显存占用。我记得第一次在自己的模型上启用融合注意力时单次前向传播的 latency 直接降低了 30% 以上。第二块是对 KV Cache 的显存优化——这是 Transformer 生成式模型最容易被忽视的显存杀手。生成长度越长KV Cache 占用越大推理时的显存占用甚至可能比模型权重本身都要高。真正让我觉得这个团队眼光毒辣的是他们推出的 DeepSpeed-MIIModel Inference for Intelligent我没有去细究缩写来源总之是一个推理加速层。MII 做了一件特别务实的事把推理优化搞成了开箱即用的 API。你只要指定一个模型名字MII 会去自动探测层结构、自动选择 kernel 融合策略、自动决定是否做量化。对于不想深入了解底层优化的团队来说这就是把系统工程师的经验直接翻译成了产品能力。评论区可能会有人问既然训练已经能跑 MoE 和 1750 亿参数模型了为什么还要专门的推理优化我可以告诉你一个真实到有点扎心的事实训练是异步的、可以容忍几秒钟的卡顿推理是同步的每一个 token 的延迟都直接体感化。到了推理阶段显存管理不再是优化器状态的问题而是离散的、动态增长的 KV Cache 和不断变化的生成长度。这完全是另一套工程难题。而前文提到的 ZeRO-Offload 在处理推理请求时也迎来了第二春。推理场景下模型权重是静态的批量请求是动态的——你可以把不用的 expert 权重卸载到 CPU根据请求的批次和路由结果动态加载。这是训练时做不到的因为训练需要连续快速地更新所有参数但推理时权重不变CPU 与 GPU 之间的搬运不会带来一致性风险所以 Offload 在推理场景下格外好用。7. 踩坑实录我在使用 DeepSpeed 时遇到的五个真实问题这部分专门写给准备在实际项目中上手 DeepSpeed 的读者。文档和论文里写的多是如何配置并跑通但要稳定地把训练跑起来绕不开那些不起眼的坑。我直接把自己踩过的问题列出来每一个都能省下你几小时到几天的排查时间。坑一优化器状态与 ZeRO Stage 的隐性不匹配最坑。DeepSpeed 允许你在配置里指定optimizer字段但如果你写的是optimizer: {type: AdamW}DeepSpeed 会使用自己的 CPUAdam 实现这和 HuggingFace 的 AdamW 在计算细节上可能不完全一样。这会导致同一个超参数下模型的收敛曲线发生偏移。我的建议是如果要复现基线实验明确指定 one-bit Adam 或者使用和 PyTorch 相同的优化器实现然后在一个小数据集上先验证 loss 曲线是否匹配。坑二ZeRO-Offload 时 CPU 内存被瞬间吃爆。启用cpu_offload: true时优化器状态会被放到 CPU 上而 CPU 内存如果只留了默认的一二十个 GB训练会在第一个优化步骤就报 OOM。我踩过一次特别糟糕的模型 40 亿参数我给了 CPU 64GB结果它在第二个 epoch 时直接卡死了——因为除了优化器状态它还往 CPU 上放了临时缓冲区。后来我学会了一个保守的估算公式CPU 内存需求约等于模型参数量 × 18 字节FP16 权重 2 FP32 梯度 4 FP32 优化器状态 1240 亿参数至少需要 72GB CPU 内存。坑三梯度累积与 ZeRO Stage 的协调问题。用gradient_accumulation_steps: 8的时候ZeRO Stage 2 的梯度分片必须等所有 micro-batch 都完成反向后才能触发一次 ReduceScatter。如果你把 accumulate 设得很大或者 micro-batch 数不能被数据并行度整除某些卡可能会提前触发通信造成死锁或数值不一致。最好的做法是保证 gradient accumulation 是 data_parallel_size 的倍数或者干脆从一开始就把 batch size 拆得足够小让 accumulation 次数保持相对少。坑四NCCL 域内带宽跑不满通常不是你配置文件的错。很多人发现同样的配置在 8 卡机上速度正常扩到 32 卡后吞吐增长反而很蹊跷。一个容易被忽略的真相是NCCL_P2P_DISABLE和NCCL_SOCKET_IFNAME这些环境变量只影响通信库和 DeepSpeed 本身关系不大。跨节点通信时如果网络不支持 RDMA远程直接内存访问AllGather 和 ReduceScatter 的吞吐会断崖式下降。检查你的节点之间有没有真正接通 RDMA比调整 DeepSpeed 配置参数重要得多。坑五断点续训和权重转换的格式兼容问题。DeepSpeed 的 checkpoint 格式和原生的 PyTorch checkpoint 不完全互通。stage3_gather_16bit_weights_on_model_save这个参数如果没设你 save 出来的会是分片权重而不是完整权重导致你想直接拿去跑评估脚本时找不到state_dict。后来我的标准做法是保存 checkpoint 时显式设置该参数为 true它会多花一点时间收集权重并且每个例如 2000 步保存一次避免权重全丢。8. ZeRO 与未来趋势十年演进没有停下的意思说到近期的演进绕不开 2023 年提出的 ZeRO。它的目标很明确解决 ZeRO Stage 3 在跨节点场景下通信开销过大的问题。普通的 ZeRO Stage 3 在每次前向和反向传播中都需要对权重做 AllGather在每次反向结束时对梯度做 ReduceScatter。当模型规模达到千亿以上时这些通信量是天文数字通信耗时甚至可能超过计算耗时。ZeRO 的核心手段之一是权重分区量化Quantized Weight Partition把权重整数化后传输传输量直接减半以上同时引入了分层式 AllGather让通信不再完全依赖全局 AllGather而是通过部分复制和分层传播来降低峰值带宽需求。这里我不过度展开原理细节简单说就是——把用通信换显存的代价进一步压缩让 ZeRO 在高带宽跨节点场景下也更具竞争力。除了 ZeRO更值得关注的是 DeepSpeed 整个项目在 2023 年后期的位置变化。随着 PyTorch 原生 FSDPFully Sharded Data Parallel的成熟DeepSpeed 独占的自动分片优势已经被抹平了不少。FSDP 在易用性上做得非常好和 PyTorch 生态的兼容性天然更顺。但 DeepSpeed 在训练侧积累的通信优化、推理侧的多模型支持、以及 Chip-to-Chip 的低延迟优化依然是很多人离不开它的理由。在我看来未来一两年的演进会集中在三个方向Sparsity 和 MoE 的深度融合。不管是学术界还是工业界都在用稀疏架构绕开稠密模型的显存瓶颈而 DeepSpeed 的 Amoeba 和 MoE 工具链已经有了先发优势。推理引擎与服务化一体化。MII 已经走出了半步未来很可能把训练完的权重直接优化并部署为可调用的推理服务把管训练和管推理彻底打通。自动化配置的普及。DeepSpeed 一直在尝试让系统自动选择最优的并行策略和 offload 策略ds_report和 autotuning 功能会越做越强。未来用户可能只需要给出模型和硬件资源DeepSpeed 自动给出配置和预期的吞吐曲线。我不太喜欢把这种趋势说得过于玄学。对于普通从业者来说最实际的建议是在 2024 年做分布式训练选型时如果你更看重工程的稳定性和社区文档的成熟度DeepSpeed 依然是首选如果你只想在 PyTorch 生态内快速实现分片并跑通实验FSDP 可能学得更快。两个生态会继续互相借鉴、互相竞争真正受益的是我们这些每天和显存、带宽、NCCL 报错搏斗的人。这十年DeepSpeed 从一个小众的优化器工具长成了大模型训练体系里绕不开的一块基石。它的演进史本质上也是大模型系统化工程从野蛮生长走向标准化的一个缩影。如果你问我做技术选型时最看重什么我会说看一个项目怎么对待用户不想改代码这件事。DeepSpeed 让我觉得舒服的地方就是多数时候它确实做到了把复杂留给自己把简单留给用户。这大概也是它能在激烈的框架竞争中活到今天的核心原因吧。