DeepSpeed训练方式

DeepSpeed 是基于 PyTorch 的深度学习优化库，支持数据并行、张量并行和流水线并行，并通过 ZeRO 分片减少训练状态的重复存储。 模型结构、训练数据与损失函数仍由训练程序定义。

数据并行分配不同样本，模型并行划分模型计算。张量并行和流水线并行是两种常见的模型并行方式。

ZeRO 则减少数据并行中重复保存的训练状态。单独使用 ZeRO，不会把网络计算切成不同层或不同矩阵块。

1. 并行方式

1.1 数据并行

数据并行（Data Parallelism，DP）让多个模型副本处理不同样本，汇总梯度后更新参数。 普通 PyTorch DDP（DistributedDataParallel）通常由每张 GPU 上的一个进程保存完整模型，各进程执行相同的网络结构。

例如，4 张 GPU 处理一个包含 32 个样本的 batch，每张卡分到 8 个样本。各卡对自己的样本做前向和反向传播，汇总梯度后更新参数。虽然处理的样本不同，各卡使用的是汇总后的梯度，更新后的模型参数保持一致。

DDP 在反向传播过程中同步梯度。每个进程分别创建普通优化器时，参数、梯度和优化器状态都会在各卡上重复保存。增加数据并行 GPU 数量能分担样本计算，但不会自动缩小单卡的完整模型状态。

1.2 模型并行

模型并行（Model Parallelism，MP）把一个模型的计算分给多张 GPU，一个样本需要这些 GPU 共同完成计算。 常见划分方式是层内的张量并行，以及按层分段的流水线并行。

    flowchart LR
    A[多 GPU 训练] --> B[数据并行 DP]
    A --> C[模型并行 MP]
    C --> D[张量并行 TP]
    C --> E[流水线并行 PP]
  

“模型并行”在不同资料里的含义可能不同。例如，有些资料把它与“流水线并行”并列，此时的“模型并行”特指张量并行。本文采用广义的模型并行，包含张量并行（TP）和流水线并行（PP）。

1.2.1 张量并行

张量并行（Tensor Parallelism，TP）把同一层中的张量和运算划分给多张 GPU。 每张卡只计算这一层的一部分结果，通过通信把各部分结果组合起来，或继续以分片形式交给后续运算。

例如，把不含偏置的线性层 $Y=XW$ 分给两张 GPU。权重按输出维度切成两块，写成 $W=[W_0,W_1]$。GPU 0 计算 $XW_0$，GPU 1 计算 $XW_1$。两张卡使用相同的输入，分别得到输出的一部分；把两部分按输出维度拼接，就是完整的 $Y$。

实际网络不必每层都拼成完整输出。相邻层可以配合分片布局，减少收集操作，但前向和反向仍需要相应的跨卡通信。

TP 能减少单卡承担的层内权重和计算，适合单层很大、单卡难以容纳的模型。它在层内频繁通信，对 GPU 互联带宽和延迟更敏感。

DeepSpeed 的 AutoTP 能根据支持的模型结构或分片规则划分张量；自定义模型需要提供可识别的划分方式。仅设置 GPU 数量，不足以让任意模型正确执行 TP。

1.2.2 流水线并行

流水线并行（Pipeline Parallelism，PP）按层把模型分成多个阶段，每个阶段放在不同 GPU 上。 前向传播把激活传给下一阶段，反向传播把激活的梯度传回上一阶段。

只把层分开放置时，一个样本仍要依次经过各阶段，GPU 会等待上游结果。PP 将一个 batch 拆成多个 micro-batch（微批次），让不同阶段同时处理不同 micro-batch。

例如，一个 8 层模型分给两张 GPU，GPU 0 负责第 1–4 层，GPU 1 负责第 5–8 层。将一个 batch 拆为 A、B、C 三个 micro-batch，只看前向计算：

时段	GPU 0：第 1–4 层	GPU 1：第 5–8 层
1	处理 A	等待 A 的激活
2	处理 B	处理 A
3	处理 C	处理 B
4	本批前向结束	处理 C

时段 2、3 中，两张卡同时计算，但处理的是不同 micro-batch。真实训练还要安排反向传播与梯度累积，具体执行顺序由流水线调度决定。

流水线启动和结束时的等待称为 流水线气泡（pipeline bubble）。各阶段耗时不均时，较快的阶段还会等待较慢的阶段，因此按层数平均分配不一定能让各卡工作量相同。

DeepSpeed 的 PipelineModule 描述层与阶段的划分，PipelineEngine 调度各 micro-batch 的前向、反向和参数更新。

1.3 混合并行

DP、TP 和 PP 可以在同一次训练中共同使用，通常称为 3D 并行。 PP 划分网络阶段，TP 分担阶段内部的层计算，DP 让多组 GPU 处理不同样本。

例如，8 张 GPU 采用 DP=2、PP=2、TP=2，可按下面的方式分工。GPU 编号仅用于说明分组：

模型副本	前半层阶段	后半层阶段	输入
副本 A	GPU 0、1 做 TP	GPU 2、3 做 TP	一组样本
副本 B	GPU 4、5 做 TP	GPU 6、7 做 TP	另一组样本

每个副本由 4 张卡共同完成模型计算；后半层阶段接收前半层传来的激活。GPU 0 与 4、1 与 5、2 与 6、3 与 7 分别组成数据并行组。每对 GPU 负责模型的同一部分，在组内汇总对应参数的梯度。

虽然总共有 8 张卡，处理不同样本的模型副本只有 2 个，因此数据并行度是 2。每个副本内部的 TP 和 PP 共同处理同一组样本。

ZeRO 在数据并行组内分片状态，能否与某种 TP、PP 实现一起使用，要看具体引擎的支持范围。3D 并行描述 GPU 的分工，不代表 AutoTP、PipelineModule 和所有 ZeRO 阶段都能直接组合。

2. 训练显存

训练时，GPU 除了保存模型参数，还需要保存梯度、优化器状态，以及反向传播要用的中间结果。

混合精度训练还可能保留 FP32 主权重：参数更新在这份精度更高的权重上完成。更新后的结果会转换成低精度权重，供前向和反向计算使用。

2.1 模型状态大小

参数、梯度和优化器状态合称模型状态。下面只估算这部分显存。

ZeRO 原论文采用 FP16＋Adam 设置分析显存：参数和梯度为 FP16，主权重与 Adam 的两个矩为 FP32。假设全部参数参与训练，每个参数对应的状态为：

后文沿用这组假设，把 FP32 主权重与两个矩合计的 12 字节归入优化器状态。设参数量为 $P$，未分片的模型状态占用就是 $16P$ 字节。

16 字节只对应这组精度与优化器假设，不包含激活和临时空间。 实际实现可能用 FP32 累积梯度，或采用不同的主权重、优化器状态精度，不能把它当成所有训练任务的固定系数。

3. ZeRO 状态分片

ZeRO（Zero Redundancy Optimizer，零冗余优化器）在数据并行进程之间分片保存训练状态。 普通 DDP 中重复保存的状态，分片后由各进程分别保留其中一部分。

ZeRO 的分片范围逐级增加：Stage 1 分优化器状态，Stage 2 加上梯度，Stage 3 加上参数。

设数据并行组内有 $N$ 个进程，每个进程使用一张 GPU，各阶段保留的模型状态如下：

“分片”描述需要保留的状态；计算和通信过程中仍会产生临时张量。$N$ 是参与这次分片的数据并行组大小，混合其他并行方式时不一定等于总 GPU 数。

3.1 ZeRO-1：分片优化器状态

Adam 更新某个参数时，只需要这个参数的梯度、主权重及对应的一阶矩和二阶矩。其中的主权重和两个矩，不必在每张卡上各存一份。

ZeRO-1 为每个进程分配一部分参数的更新任务。梯度完成同步后，各进程更新自己负责的参数片段，并交换更新后的参数，使每张卡重新拥有完整的低精度参数。

例如，用两个进程训练同一个模型，把参数按编号分成前后两半。进程 0 保存前半部分的优化器状态并负责更新这部分参数，进程 1 负责后半部分。交换更新结果后，两边都拿到完整的新参数，继续处理各自的样本。

按照前面的精度假设，优化器状态占 $12P$ 字节，是模型状态中的大头。ZeRO-1 保留完整参数和梯度，仍能明显减少重复存储。

3.2 ZeRO-2：分片梯度

ZeRO-1 已经把更新任务分给不同进程。每个进程更新参数时，只需要自己负责的那部分梯度。因此，ZeRO-2 让各进程只保留同步后属于自己的梯度分片。

两种通信操作的区别是：

• All-Reduce：汇总各进程的张量，每个进程都拿到完整的汇总结果。 普通 DDP 主要通过它同步梯度。
• Reduce-Scatter：汇总各进程的张量，并把汇总结果切成分片，每个进程只拿到对应的一片。

ZeRO-2 的梯度归约可以按 Reduce-Scatter 的方式理解：每个进程接收自己负责的梯度分片，用对应的优化器状态更新参数。各进程交换更新后的参数片段，重新得到完整参数。 实际底层通信会受实现和配置影响。

梯度分片不意味着某个参数只使用一张卡上的样本学习。每个梯度分片仍汇总了数据并行组内各进程的贡献。

3.3 ZeRO-3：分片参数

ZeRO-2 仍要求每张卡保存完整的低精度参数。参数本身很大时，这部分显存仍可能放不下。

ZeRO-3 将参数也分片，计算某个模块时临时收集所需的完整参数。计算后，这些临时参数会被释放，或暂时保留以便复用。 All-Gather 是其中的收集操作：各进程交换已有的分片，得到当前计算所需的完整参数。

以下展示一次模块计算中参数的状态变化，省略预取和缓存：

    flowchart LR
    A[各卡持有分片] -->|All-Gather| B[收集模块参数]
    B --> C[处理各自样本]
    C --> D[释放临时参数]
    D --> A
  

前向传播时，各进程按模块收集参数。反向计算如果需要已经释放的参数，就重新收集。计算出的梯度经汇总、分片后，由负责该分片的进程更新对应参数。更新后的参数继续以分片形式保存。

ZeRO-3 改变参数的保存方式；在没有叠加模型并行时，每张卡仍对自己的样本执行完整网络的前向和反向计算。

ZeRO-3 会增加计算期间的参数通信。提前传输下一模块的参数（预取）、让通信与计算同时执行，或保留即将复用的参数（缓存），都有机会减少等待。其中，预取和缓存会让更多参数同时留在显存中。

3.4 显存估算

沿用第 2 节的精度假设，只计算模型状态，均匀分片时每张 GPU 的理论状态量为：

$M$ 的单位是字节。ZeRO-1 的 $4P$ 对应完整参数和梯度，ZeRO-2 的 $2P$ 对应完整参数；其余需要保留的状态由 $N$ 个进程分担。

例如，假设模型有 $P=7\times10^9$ 个可训练参数，使用 8 张 GPU 做纯数据并行，不做 CPU 或 NVMe 卸载，并沿用 FP16＋Adam 的每参数 16 字节假设。按 $1\,\mathrm{GB}=10^9$ 字节计算，每张卡的理论模型状态量为：

• 普通 DDP：$16P=112\,\mathrm{GB}$。
• ZeRO-1：$4P+12P/8=38.5\,\mathrm{GB}$。
• ZeRO-2：$2P+14P/8=26.25\,\mathrm{GB}$。
• ZeRO-3：$16P/8=14\,\mathrm{GB}$。

这里的 14 GB 不代表 16 GB 显卡足够训练该模型。激活、梯度精度差异、参数临时收集和通信缓冲区都没有计入。

判断模型能否放入显存时，还要考虑最大模块的参数量、激活和计算所需的临时空间。DeepSpeed 的内存估算工具只覆盖其中一部分，显存峰值仍要在实际训练配置下测量。

上面的 $P$ 是纯数据并行中的完整模型参数量。叠加 TP 或 PP 后，各卡负责的模型部分已经改变，需要按实际持有的参数重新计算，不能只修改 $N$。

3.5 ZeRO 与 FSDP

FSDP（Fully Sharded Data Parallel）是 PyTorch 的分片数据并行实现。 在全分片方式下，它与 ZeRO-3 都分片保存参数、梯度和优化器状态，并在计算时收集参数。

FSDP 和 ZeRO-3 都没有把完整参数长期保留在每张卡上，但计算当前模块时仍需取得所需参数。TP 则让各卡只执行该层的一部分运算，这是状态分片与计算划分的区别。

FSDP 与 DeepSpeed 的接口、分片调度和检查点格式不同，迁移时需要调整对应的训练与保存逻辑；相近的分片机制也不意味着相同的训练速度。

4. 显存优化

4.1 混合精度

混合精度让不同计算或状态采用不同的浮点精度。 较低精度通常能减少存储与传输量，在支持相应运算的硬件上加快计算；需要更高精度的累积和更新仍可保留 FP32。

FP16 的数值范围较小，训练常配合 loss scaling：放大损失和相应梯度，减少小梯度下溢，更新前还原梯度尺度。BF16 的指数范围与 FP32 相同，通常不需要 loss scaling，但有效精度仍低于 FP32。

DeepSpeed 提供 fp16、bf16 等配置入口。这两个原生模式不能同时开启，选择 BF16 还需要硬件支持。启用低精度计算，不等于参数、梯度和优化器状态全部采用同一种低精度。

4.2 梯度累积

梯度累积将一次参数更新所需的样本拆成多个 micro-batch，分次前向和反向，累计梯度后执行一次更新。 在纯数据并行中，micro-batch 就是一张 GPU 一次前向和反向处理的样本集合。

在每个进程的 micro-batch 大小相同、每次累积步数固定的纯数据并行训练中：

$B_{\mathrm{global}}$ 是一次参数更新涉及的样本数，也就是有效 batch；$B_{\mathrm{micro}}$ 是单卡 micro-batch 大小，$A$ 是累积步数，$N$ 是数据并行进程数。

例如，8 张 GPU 的单卡 micro-batch 为 2，累积 4 次梯度，每次参数更新就涉及 $2\times4\times8=64$ 个样本。把单卡 micro-batch 降为 1、累积次数增为 8，仍是 64 个样本，但每次前向需要同时保存的激活通常更少。

固定有效 batch 时，减小 micro-batch 才是降低激活峰值的关键；只增加累积次数不会减少原有的参数和优化器状态。 累积期间仍要正常完成各 micro-batch 的反向传播，不能把所有前向计算图一直保留到最后。

叠加 TP、PP 时，公式中的 $N$ 仍是数据并行度。 一个模型副本内部的多张卡共同处理同一 micro-batch，不额外增加样本数。

例如，沿用第 1.3 节的 8 卡混合并行配置，数据并行副本只有 2 个。若每个副本的 micro-batch 为 2，累积 4 次，则每次更新涉及 $2\times4\times2=16$ 个样本，不能按总卡数算成 64 个。

这三个 batch 参数对应的 DeepSpeed 配置是：

• 有效 batch $B_{\mathrm{global}}$：train_batch_size。
• 单卡 micro-batch 大小 $B_{\mathrm{micro}}$：train_micro_batch_size_per_gpu。
• 累积步数 $A$：gradient_accumulation_steps。

相同的样本数不一定对应相同的 token 数：序列越长，token 越多。因此，比较训练吞吐时还需要记录输入长度。

4.3 激活检查点

激活检查点（Activation Checkpointing）减少前向保存的中间张量，反向传播时根据保留的输入重新计算所需结果，以额外计算换显存。 它不减少 Adam 的状态。

例如，把连续几个 Transformer 层划为一个重算区间。前向保留区间输入，省去部分内部激活的长期保存；反向需要这些激活时，重新执行相应的前向计算，再求梯度。

模型或训练框架需要指定哪些模块参与重算，并接入 DeepSpeed 等库提供的激活检查点接口。仅开启 ZeRO，不会自动为任意模型的所有层启用激活重算。

激活检查点服务于一次训练迭代内的反向计算；保存到磁盘、用于恢复训练的 checkpoint 是另一件事。

4.4 CPU 与 NVMe 卸载

Offload 把部分训练状态移出 GPU 显存，放进 CPU 内存或 NVMe 存储，需要时再传输到计算设备。 分片减少各进程重复保存的状态，卸载则改变状态所在的存储位置。

ZeRO-Offload 将优化器状态的存储和参数更新计算转移到 CPU，DeepSpeedCPUAdam 用于加快 CPU 上的 Adam 更新。ZeRO-Infinity 基于 ZeRO-3，进一步利用 CPU 内存和 NVMe 卸载模型状态。

GPU 显存占用降低后，主机内存容量、CPU 更新速度、CPU 与 GPU 之间的传输带宽会影响训练速度。NVMe 卸载还增加存储读写开销。卸载能扩大可训练模型的规模，但不保证训练更快。

5. 训练接口与检查点

普通 DeepSpeedEngine 的训练接口如下：

engine 是表中给引擎对象起的变量名。采用默认的梯度累积管理方式时，每个 micro-batch 都调用 step()，但只有到达累积边界才更新参数。

流水线训练由 PipelineEngine.train_batch() 安排多个 micro-batch 的计算，不能直接照搬表中的逐 micro-batch 调用流程。原生流水线引擎能配合 ZeRO-1，不能直接叠加 ZeRO-2 或 ZeRO-3。

配置文件中的 zero_optimization.stage 决定 ZeRO 阶段，精度、batch 和卸载也有对应配置。TP 需要模型的张量划分规则，PP 需要层与阶段的划分，单独修改 ZeRO 阶段不会启用它们。

恢复训练需要模型、优化器及相关训练状态；只导出模型权重无法恢复原有的优化器状态。 DeepSpeed 提供 checkpoint 保存和恢复接口，以及将 ZeRO checkpoint 合并为普通 FP32 权重的工具。使用外部训练框架时，要确认保存的是完整训练状态还是仅模型权重。

6. 并行与显存配置

配置选择取决于显存占用的来源和 GPU 的通信条件：

ZeRO-3 仍要求当前模块的参数和计算空间能放入单卡。 如果这一部分就超出显存，仅增加数据并行卡数仍可能无法训练，需要缩小单次计算规模或采用 TP 等方法。

比较吞吐时，应固定模型、输入长度、有效 batch 和精度，观察峰值显存、每秒有效 token 数及通信等待。跨机器训练还要记录节点间网络带宽。同样的 ZeRO 配置，在不同互联条件下可能表现不同。

ZeRO-3 增加参数收集，TP 增加层内通信，PP 受阶段耗时和流水线气泡影响。判断是否加速，需要比较相同训练任务的实际吞吐，不能仅看 GPU 数量或 ZeRO 阶段。