（本博客为Datawhale的baseLLM开源学习项目的学习笔记）
一、DeepSpeed全家桶一览
1.DeepSpeed概述
我们在前面已经通过QLoRA、PEFT、RLHF等技术，体验了在单机单卡甚至消费级显卡上完成大模型微调的可能性。但如果目标从"微调一个7B模型"升级为从零预训练、全量微调甚至训练万亿参数模型，仅靠量化与LoRA显然还不够，此时就需要引入微软开源的PyTorch分布式训练与推理优化库DeepSpeed。它的目标是让开发者在相同硬件条件下用更少的显存训练更大的模型、训练更快、扩展更稳，并提供从训练到推理的一整套工程化组件。它与Hugging Face的关系可以简单理解为Hugging Face更偏向"模型与数据"的生态（transformers、datasets等）而DeepSpeed更偏向"系统与算力"的生态（显存优化、并行策略、通信调度等）。
在微软的官方介绍中，DeepSpeed的主要子模块包括：
**DeepSpeed Training：**本节的主角，包含ZeRO、Offload、Infinity、3D并行等训练技术；
**DeepSpeed Inference/Compression：**针对推理和压缩（量化、剪枝、KV Cache优化等）的加速库；
**DeepSpeed-MoE：**面向MoE大模型的高效路由与通信实现；
**DeepSpeed-Chat/RLHF：**支持大规模RLHF、DeepSpeed-Chat等对话模型训练；
**DeepSpeed for Science：**将这些系统能力扩展到科学计算场景。
日常工作中，最常见的两个使用场景是当我们想在单机多卡或多机多卡环境下进行预训练或全量微调大模型时，通常会重点使用DeepSpeed Training + ZeRO。而当我们需要在服务器上部署一个高吞吐的推理服务时，则会更多地结合DeepSpeed Inference与各类量化/压缩技术一起使用。
2.DeepSpeed Training要解决的三大挑战
DeepSpeed Training希望解决的是这样一个组合问题。一是显存不够 ，单张GPU就算有80GB显存，也远远不下万亿参数的模型训练所需状态（权重、梯度、优化器状态、激活值）；二是加速比不理想 ，单纯堆GPU数量，在数据并行下很快就会遇到通信瓶颈，难以获得接近线性的加速；三是硬件异构，显示集群中同时存在GPU显存、CPU内存、NVMe SSD甚至远程存储，如何让这些异构资源一起"干活"，而不是让昂贵的GPU无限空等。围绕这三个问题，DeepSpeed提出了以ZeRO（ZeRo Redundancy Optimizer）为核心的一系列技术（ZeRO、ZeRO-Offload、ZeRO-Infinity），再叠加数据并行、模型并行、流水线并行构成完整的3D并行框架。
二、显存开销与ZeRO系列内存优化
前面我们已经学会了如何估算推理阶段的显存开销，知道"只加载权重并做前向传播"大概需要多少显存。这里我们换一个视角，聚焦在训练相对于推理多出来的那一部分显存。相比只做推理或只训练极少量LoRA参数的QLoRA微调，全量预训练或全量微调需要为每一个参数额外维护梯度和优化器状态，而这正是ZeRO要重点优化的部分。为了更直观地理解这一点，我们先以一个7B模型为例，看看在不做任何内存优化时，全量训练大致需要多少显存。
1.Qwen2.5-7B全量训练需要多少显存
假设我们想要对一个7B模型（约70亿参数）进行全量预训练或微调，并采用常见地Adam优化器和混合精度训练策略。
（1）模型权重
存储精度：通常采用FP16/BF16（2Bytes）；
显存占用：
（2）梯度
反向传播中同样以16-bit精度存储；
显存占用与参数同量级：
（3）优化器状态--以Adam为例
在优化器内部，一般会将参数提升为FP32（4Bytes）精度来计算更新；
对每个参数，Adam至少需要存储：
一个FP32精度的参数拷贝（Master Weights）；
一阶动量m；
二阶动量v。
每个状态都是4Bytes，因此单参数的优化器状态一共约12Bytes：
对于7B参数模型：
（4）总显存开销（不含激活与碎片）
把上述三部分加起来：
再考虑到中间激活、临时张量、显存碎片等，实际训练中的峰值显存往往会来到120GB甚至更高。很多资料中常见的结论，一个7B模型如果想全量训练，大概需要120GB显存，就是由这一组估算推导出来的。
为了便于后续讨论，可以把"每个参数对应的字节数"抽象写成：
也就是说，在经典数据并行+Adam的设定下，每个可训练参数平均需要约16Bytes显存（不含激活）。这就是DeepSpeed想要"动刀"的地方。如果我们能把这16Bytes中的冗余复制部分打散分摊到多张卡上，每张卡就不需要承担完整的16Bytes开销，从而在固定总GPU数量下训练更大的模型
2.Zero Redundancy Optimizer的分治思路
回到刚才的结论，在经典数据并行+Adam的设定下，每个参数都会对应一整套模型状态。在Data Parallel（数据并行）中，会用N张GPU，每张GPU上都存一份完整的模型状态副本，每次训练时由不同GPU处理不同的数据子集（batch切分）计算结束后再通过all-reduce汇总梯度并统一更新参数。这样一来，那一整套模型状态会在每张卡上各复制一份，总共变成N份，这部分跨设备的重复存储就是所谓的"冗余"。ZeRO的名字"Zero Redundancy Optimizer"正是针对这一痛点而来，它的核心思路是将这些冗余的模型状态在不同GPU之间分片（Shard），从"人人持有一整份"变成"大家一起拼一份"，从而降低单卡显存压力。
2.1 ZeRO的三个Stage在分谁
ZeRO将显存优化分为三个递进的阶段，关键是逐步把优化器状态（Optimizer States） 、梯度（Gradients） 和**模型参数（Parameters）**从"每卡完整复制"变为"多卡分片存储"。
在经典数据并行下，每张GPU需要维护：
其中，
表示在经典数据并行（Data Parallel, DP）设定下单张GPU上用于存储模型状态的开销，而OS、G、P分别代表优化器状态、梯度和参数，这里只关注它们的相对规模，所以用
表示与这三部分之和成正比。
这个过程中ZeRO分别在三个Stage中逐步把这三类状态打散到多张GPU上：
（1）ZeRO-1（Optimizer States Sharding）
在这一阶段，ZeRO仅对优化器状态（OS）进行分片，而参数（P）与梯度（G）仍然在每张GPU上完整保存。此时，单卡显存开销近似变为：
相当于大家依然各自持有完整的模型和梯度，但优化器状态（如Adam的动量）被切分，每张卡只负责维护其中一小块。由于优化器状态在总显存中占比往往最大（约75%），仅这一步就能带来显著的显存节省。这也是为什么在当前的大模型训练实践中，ZeRO-1往往被视为"最低配置"或默认开启的基线功能。
（2）ZeRO-2（OS + G分片）
在ZeRO-1的基础上，ZeRO-2进一步对梯度进行了分片，此时只有参数仍完整复制在每张卡上。单卡显存进一步降低为：
这通常能带来约8-16倍的显存节省（视DP并行度而定），使得单卡能训练更大的模型。
（3）ZeRO-3（OS + G + P全分片）
这是"火力全开"的终极形态，对模型参数本身也进行了分片。此时每张GPU只保存整个模型参数的一部分，单卡显存变为：
理论上，随着GPU数量的增加，单卡显存占用可以被压缩到原来的1/DP量级，这就是训练万亿参数模型的关键所在。
如下图所示，展示了Baseline与ZeRO三个阶段在显存占用上的对比。随着ZeRO阶段的深入，单卡显存占用（绿色代表优化器状态，橙色代表梯度，蓝色代表参数）被逐步"削减"和分摊。
如果把之前说的"16Bytes/Param"的分解代入（其中OS = 12Bytes, G = 2Bytes, P = 2Bytes），可以用一个更直观的形式表示：
为了更直观地感受这三个阶段地威力，我们可以参考ZeRO论文中地实测数据（下表）表中展示了不同参数规模的模型，在不同GPU数量下，开启ZeRO各阶段后的单卡显存占用（单位：GB）：
能够看到如果不使用ZeRO（相当于第一行DP=1的情况），训练一个1T参数模型需要单卡16TB显存，这在现有硬件上是不可能的。但如果开启ZeRO-3并在1024张GPU上并行（最后一行最右侧）单卡显存仅需15.6GB这意味着用现有的16GB/32GB显卡集群就可以训练万亿参数模型。
2.2 通信开销与"线性加速比"的平衡
在分片带来显存红利的同时，也会引入额外的通信开销。ZeRO论文中对通信量（Communication Volume）进行了详细分析（假设模型参数量为
）：
ZeRO-1/ZeRO-2： 通信量为2，与标准数据并行同阶。直观理解是把标准DP的all-reduce用reduce-scatter + all-gather这类等价组合来实现后，总通信量量级不变，但通信"形态"发生了变化（例如ZeRO-2需要配合分片优化器在step内做参数分片的聚合/同步）
ZeRO-3： 通信量为3，约为标准数据并行的1.5倍。这是因为在前向/反向传播中需要为计算临时聚合当前层所需参数分片（可理解为额外的参数all-gather）再叠加梯度的reduce-scatter与更新后参数的同步通信。虽然通信量上升，但换来的是参数级别的全分片，使得单卡显存占用能随GPU数量近似按1/DP下降。
总的来说，ZeRO系列论文和DeepSpeed的工程实现，核心就是在显存开销、通信带宽与计算时间 三者之间找到一个工程上可接受、又足够通用的平衡点。由于ZeRO节省了大量显存，往往允许使用更大的Batch Size，可以显著提升计算的算术强度，在部分场景下甚至能观察到**超线性加速（Super-Linear Speedup）**的效果。
3.ZeRO-R
在实践中，即使把参数、梯度和优化器状态都进行了分片优化，训练过程中的显存峰值仍然可能被其他因素"偷走"。例如，中间激活值的缓存、反向传播中临时张量的创建与释放、显存碎片（Fragmentation）导致的不可用空间，以及不均匀的layer配置导致的部分GPU负载不均，都可能成为新的显存瓶颈。
为此，论文中提出了一个补充模块ZeRO-R，主要从三个维度进一步优化显存使用：
分区激活检查点（Partitioned Activation Checkpointing, 
）
结合Activation Checkpointing，但在模型并行（MP）中，ZeRO进一步删除了激活值的复制冗余；
仅保留分片的激活检查点，需要时通过all-gather重建。这使得激活显存随MP并行度线性降低。
恒定大小缓冲区（Constant Size Buffers）
将很多临时张量的分配重定向到一个或少数几个可复用的大型buffer中；
避免频繁的小块分配与释放造成的显存碎片。
显存碎片整理（Memory Defragmentation）
对显存中的可用区域做更加智能的分配策略；
在长时间训练过程中保持可用显存的有效比例。
这些优化看起来只是"工程小技巧"，却往往决定了：你是刚好能跑起来，还是训练到一半OOM崩溃
4.ZeRO-Offload
4.1 为什么需要Offload
前面的ZeRO假设所有模型状态（无论是否分片）都存放在GPU显存中。但在许多现实环境下，GPU显存资源往往非常有限（比如每卡仅16GB-24GB），而CPU内存则相对"富裕"（几十GB甚至上百GB）如果能把一部分状态"搬"到CPU内存上，哪怕训练速度稍慢一点，也能让普通用户有机会训练参数量大得多的模型。
ZeRO-Offload的思路是把一部分原本需要占用CPU显存的模型状态（最典型的是FP32 master weights与优化器状态，以及在某些配置下的梯度缓冲区），卸载（Offload）到CPU内存上，由CPU负责更新；GPU侧则尽可能聚焦前向/反向的计算。论文从训练数据流与通信量角度分析了这种"GPU计算前向/反向+CPU计算参数更新 "划分在多种约束下的合理性。下图展示了其核心架构：GPU负责计算密集型的前向与反向传播（FWD-BWD Super Node）CPU负责显存密集型的参数更新（Update Super Node）。图中M表示模型参数量，对应箭头上的2M/4M可理解为一次step内在CPU与GPU之间需要交换的数据量级（常见的拆解是：梯度约2M（FP16）参数回传约2M（FP16）合计约4M；而FP32 master与优化器状态主要驻留在CPU内存中）
4.2 单机单卡下的典型数据流与优化
在单卡+ZeRO-Offload的设定下，一次训练step的简化流程大致如下：
（1）GPU上前反向传播，必要时将梯度/分片梯度Offload到CPU：
GPU上保留用于计算的FP16/BF16参数副本，完成前向计算得到激活与loss，并进行反向传播得到梯度。随后根据配置（例如是否做梯度分片、是否将梯度缓冲转移到CPU）通过PCle/NVLink将梯度（或梯度分片）拷贝到CPU内存，以释放GPU显存压力并为CPU侧更新做准备。
（2）CPU上执行优化器更新（更新FP32 master与优化器状态）：
CPU内存中保存着FP32精度的master weights以及Adam的m、v等状态。CPU利用刚刚拷贝来的梯度更新这些状态，得到新的参数值（此时仍在CPU端）。为了降低CPU侧优化器成为瓶颈的风险，ZeRO-Offload会采用高度优化的CPU Adam实现（例如利用SIMD指令与多线程并行）。同时，为了进一步掩盖CPU计算与CPU-GPU交换的延迟，论文还提出了单步延迟更新（One-Step Delayed Parameter Update，DPU） 机制，允许CPU的参数更新与GPU下一步的前向/反向计算在一定程度上重叠执行。下图用"计算流（Computation stream）"与"交换流（Swapping stream）"展示了这种overlap的直觉：Step i的梯度产生后，CPU开始更新参数；GPU则可以更早进入Step i+1的计算阶段，从而减少等待。这种重叠执行有效地掩盖了CPU计算和CPU-GPU通信带来的延迟，保证了训练的高吞吐量。虽然GPU可能会使用旧一步的参数进行计算，但实验证明这几乎不影响收敛性。
（3）将更新后的参数（FP16计算副本）回传到GPU：
CPU完成更新后，将用于计算的FP16/BF16参数副本同步回GPU，进入下一个训练step的循环（实际系统里通常会分层/分bucket同步，并尽量与计算重叠）
通过这种设计，ZeRO-Offload在单卡V100（32GB） 的设置下可将可训练模型规模提升到**10B量级（论文示例中最高约13B）**并在一定条件下保持较高吞吐（数量级可达数十TFLOPS）。不过在多卡训练场景下，ZeRO-Offload通常不会简单地让每张卡各自与CPU做大规模同步（这会更容易触及PCle/主机内存带宽瓶颈）而是结合ZeRO-2的梯度分片机制来控制交换量，先在GPU端对梯度做Reduce-Scatter，将N张卡上的梯度切分成N份；随后，每张GPU只需将自己负责的那1/N份梯度Offload到CPU；接着，CPU并行更新对应的参数分片；最后更新后的参数分片回传到各自GPU，并通过All-Gather同步到需要完整参数副本的设备上。得益于这种"先分片、再offload"的设计，CPU-GPU的交换开销通常不会随卡数线性膨胀，更有利于扩展到更大规模集群。
5.ZeRO-Infinity
5.1 突破显存与带宽墙
如果说ZeRO-Offload把CPU内存也纳入了训练资源，那么Ren等人提出的ZeRO-Infinity则更进一步，把NVMe SSD甚至远程存储也纳入了统一的"内存池"视角。它的思路是构建一个多层次的异构内存体系，GPU HBM显存作为Fast层，容量最小但带宽最高（几百GB/s）仅用于存放当前正在参与计算的参数shard与激活；CPU DRAM作为Medium层，容量较大但带宽较低，存放近期会被访问的参数、优化器状态与激活片段；NVMe/SSD则作为Slow层，容量最大但带宽延迟最差，用于存放远期才会被访问或备用的模型状态。Infinity Offload Engine能够自动管理这三层存储之间的数据流动，实现全量Offload。所以，只要NVMe硬盘够大，理论上就可以存放任意规模的模型状态，而不仅仅局限于CPU内存大小。
不过，将数据放在NVMe上最大的挑战是带宽（PCle 3.0/4.0远慢于HBM）。ZeRO-Infinity提出了Bandwidth-Centric Partitioning（以带宽为中心的切分）策略。传统的Offload（如ZeRO-Offload）通常受限于单张卡的PCle带宽，而ZeRO-Infinity利用All-Gather通信模式，让所有GPU同时从CPU/NVMe拉取不同的数据分片。其效果是有效加载带宽随GPU数量线性增加。例如在64张GPU的集群上，它能利用64条PCle通道的聚合带宽，从而掩盖NVMe读写慢的问题。
下图展示了带宽为中心的分区策略，模型状态被切分并存放在慢速存储（CPU+NVMe）中。与传统的让单个GPU负责所有数据传输不同，ZeRO-Infinity让每张GPU（
到
）并行地通过All-Gather操作仅拉取自己负责的那一小部分数据（如
）这样一来，整个系统的有效带宽就变成了所有GPU PCle带宽的总和，极大地提升了从慢速存储加载数据的效率。
5.2 突破单层限制
即使模型总状态能存下，超大模型（如万亿参数）的单个层 也可能大到无法放入单张GPU的显存中。ZeRO-Infinity引入了Memory-Centric Tiling（以内存为中心的切片） 把一个巨大的Operator（例如超大矩阵乘法）拆解为一系列更小的Tiles，每次只加载一个Tile到GPU参与计算，算完释放，再加载下一个。需要注意的是，这类tiling能在一些场景下降低对张量并行（TP）的刚性依赖，但并不意味着"完全替代模型并行"。工程上仍需结合算子形态、通信带宽与吞吐目标综合选择并行策略。
得益于这些技术，ZeRO-Infinity实现了自动化模型切分，即在模型初始化阶段（init）就自动进行切分和Offload，防止初始化时就OOM。同时，它还展现了惊人的扩展性，论文展示了在单台DGX-2节点（16张V100）上即可微调1万亿参数的模型；在512张GPU上可训练32万亿参数的模型。从设计目标来看，ZeRO-Infinity甚至将视野投向了支撑百万亿参数规模的超大模型训练。如下图所示，在相同硬件条件下（如512张V100GPU）ZeRO-Infinity能支持的模型规模（32T）比最先进的3D Parallelism（0.64T）高出50倍。
三、DeepSpeed中的并行策略与生态
1.并行策略
ZeRO主要解决的是"模型状态如何分片"这个问题，但要在大规模集群上高效训练，还需要与其他并行策略配合使用，并理解它在更大生态中的位置。常见的四类并行是：
（1）**数据并行（DP）：**把同一个模型复制到多张GPU上，让不同GPU并行处理不同的数据子集。这种方式实现简单，通信模式清晰，且相对容易扩展到几十甚至上百张卡。不过，它的缺点是每张卡都需要存储一份完整的模型副本，对于大模型而言显存开销巨大。这也是ZeRO的作用所在--通过对OS/G/P的分片，打破"数据并行必须每卡存储完整模型状态"的限制。
（2）模型并行（MP）： 将同一个模型的计算与参数拆分到不同GPU上，以突破单卡显存限制。工程实践中更常见的是按张量维度拆分（张量并行/TP） ：例如在单层内将大矩阵按列或行切分到多卡上，由多张卡协作完成同一个算子的计算。它的优点是能在不改变整体网络结构的情况下把"单层算子"做细粒度拆分；缺点是通信更频繁，对跨节点带宽更敏感（常见通信包括all-gather/reduce-scatter等）需要区分的是，ZeRO-3虽然通过参数切分大幅降低了显存占用，但它本质上仍属于数据并行的优化（计算时不拆分算子）而DeepSpeed支持与Megatron-LM等张量并行库结合，实现的是算子内部的细粒度拆分与并行计算（即多张卡协作完成一个算子的计算）。
（3）流水线并行（PP）： 将模型不同层段按"流水线工序"分配到不同GPU，例如48层Transformer，GPU0负责1-12层，GPU负责第13-24层，以此类推。如果只是做朴素的"按层切分"但不采用micro-batch，那么后续GPU往往必须等待前一段GPU产出激活才能继续计算，容易出现较长的空转时间；而流水线并行通过**Micro-batch（微批次）**将一个大Batch切分成多个小块，使得各段GPU能交错处理不同micro-batch，尽量填满计算时间片，并常与Activation Checkpointing配合以降低激活显存。实践中的难点是需要合理设计micro-batch大小与流水线深度，以减少流水线启动与结束阶段因设备空闲等待而产生"气泡"。
（4）3D并行（3D Parallelism）： 在实际的超大规模训练中，往往会同时使用数据并行（DP）、张量并行（TP）和流水线并行（PP） （常写作（DP/times TP/times PP））DeepSpeed与Megatron-LM等框架联合，在BLOOM等超大开源模型中成功实践了这种3D并行策略。简单来说就是DP解决"更多数据"的问题 ，TP/PP解决"更大模型"的问题 ，ZeRO+Offload/Infinity解决"显存不够&内存层次异构"的问题。
2.DeepSpeed生态中的其他模块
