一、核心挑战：模型太大，单卡显存装不下！
-
许多人可能会认为：训练大模型，只要显卡（GPU）数量足够多，就能水到渠成。但现实并非如此。真正制约超大规模模型训练的瓶颈，往往是单张GPU卡有限的显存容量。
-
举个例子：混合精度训练Llama 3.1-1B 这样的模型时，一张 GPU（如V100S 32GB）的显存就可能被撑满。即使使用显存更大的显卡（如A100 40/80GB），一个3B参数的模型也可能超出其单卡容纳能力。
-
因此，为了充分利用海量GPU的计算能力并克服单卡显存限制，我们必须借助多种并行技术，将模型的巨大显存需求和计算负载分担给不同的GPU。
二、并行方式与并行组：DP和DP组、TP和TP组、PP和PP组
1. 几种并行的基本概念
-
张量并行 (Tensor Parallelism, TP): 将单层神经网络（如Transformer的一层）内部的计算（如矩阵乘法）按照特定维度（行/列）拆分到多个GPU上协同完成。每个GPU只负责计算该层的一部分。
-
流水线并行 (Pipeline Parallelism, PP): 将整个深度神经网络按层切割成多个连续阶段 (Stages)。每个阶段包含若干层，被分配到不同的GPU组 (称为PP组或流水线组) 上。数据像流水线上的产品一样，依次经过这些阶段进行处理。
-
数据并行 (Data Parallelism, DP): 在多个GPU上复制整个模型的副本。每批训练数据被拆分成更小的微批次 (Micro-batches)，同时分发给这些副本进行独立并行计算，最后通过通信来同步模型更新。这主要用于增大每次训练迭代的数据量（Batch Size）
2. 并行组的划分
在实际训练集群中，GPU会根据所采用的技术被组织成不同的组 (Group)：
-
张量并行组 (TP Group): 一组协同完成同一层Transformer计算的GPU。它们通常位于同一个服务器内，利用高速卡间通信（如NVLink）进行频繁交互。(例如：一台服务器内的Rank 0-3)。
-
流水线并行组 (PP Group / Pipeline Group): 一组协同完成整个模型所有层计算的GPU。其中包含多个阶段（Stages），每个阶段本身就是一个TP组。同一个PP组的GPU共同存有一个完整的模型。(例如：跨服务器的Rank 0-7)。
-
数据并行组 (DP Group): 一组存储着相同模型分区的GPU。这里所说的“相同模型分区”是指：在DP组内的所有GPU上，其存储的模型部分（经过TP和PP切分后的部分）是完全一致的副本。DP组主要用于承载更大的数据Batch，并协调副本间的参数更新。(例如：不同服务器上存储相同切片的GPU，如Rank 3 和 Rank 11)。
3. 作为一块GPU，你需要知道什么？
在训练框架中，每一块参与训练的GPU，通过调用框架接口，通常能获得以下关键信息：
-
自己的
RANK号（全局唯一标识符）。 -
自己所在的
DP Group。 -
自己所在的
TP Group。 -
自己所在的
PP Group (Pipeline Group)。 -
自己在PP组内所处的阶段（Stage）。
有了这些信息，你就可以参与到大模型的训练中去了。
三、DP如何起到降低单卡显存压力的作用
-
如上所述，当使用TP和PP时，模型本身被分割为了不同的部分，所以可以减轻单卡的压力，但是使用DP的时候，每一组DP的副本都保留了一份完整的模型参数，好像并不能减轻单卡的显存压力？
-
DP是如何最终帮助我们降低单卡显存压力的呢？答案是引入ZeRO (Zero Redundancy Optimizer) 技术。
ZeRO的核心思想是：在DP组内部打破完全复制的模式：
-
让DP组中的不同GPU成员，分别只保存部分优化器状态 (ZeRO-1)、部分梯度 (ZeRO-2)、甚至部分模型参数 (ZeRO-3)。
-
计算需要时才通信： 当某块GPU需要用到不属于自己管理的模型参数或梯度进行前向/反向计算或参数更新时，才通过高效的DP组内通信（All-Gather）将它临时获取过来，计算完毕后再丢弃。
简言之，ZeRO使得DP组内的GPU能够相互借用显存资源，显存开销由组内成员分担，从而显著降低了单卡所需的峰值显存。
(示意图，不准确)
四、假如你是一块GPU，在3D并行**（DP+TP+PP）**中，你做了什么
-
在这一单元，我们假定自己是一块编号为i的GPU，在我们已经知道自己所在的DP组、TP组和PP组以后，我们该如何与其他显卡协同，共同通过三种并行方式的协同，完成超大规模模型的训练。
-
我们假设这块卡是随机一张卡，不妨假定他处于PP组中的一个中间阶段（即既有前驱阶段，也有后继阶段）。
-
ZeRO-2 以及 ZeRO-3 技术与PP并行不兼容，因此当使用ZeRO-2和ZeRO-3时，只考虑张量并行，原因将在第五章节解释。
1. 使用TP和PP并行，每个DP组都保留完整的参数和优化器(ZeRO-0 DP)
-
前向传播：
-
-
通过PP组通信，拿到前驱阶段GPU的激活值
-
通过TP组通信， 在TP组内与其他GPU协同，基于接收到的激活值完成本阶段（即你负责的若干层）的前向计算。
-
通过PP组通信，将当前部分层的激活值传输给后继阶段的GPU
-
-
反向传播：
-
-
通过PP组通信，拿到后继阶段GPU的梯度
-
通过TP组通信，在TP组内与其他GPU协同，基于接收到的梯度和本地的激活值缓存，完成本阶段的反向计算，计算出本阶段的局部梯度。
-
通过PP组通信，将当前阶段的局部梯度传输给前驱阶段的GPU
-
通过DP组通信，将所有DP组计算得到的梯度通过All-Reduce进行全局平均，每张卡都有完整的全局平均梯度
-
-
参数更新：
-
- 无需通信，你现在有模型并行分割出的子模型的完整模型和完整全局平均梯度，直接用来更新本地的参数即可
2. 使用ZeRO-1 DP(DP组的每一个成员保存一部分优化器状态)
-
前向传播：(没有变化)
-
-
通过PP组通信，拿到前驱阶段GPU的激活值
-
通过TP组通信， 在TP组内与其他GPU协同，基于接收到的激活值完成本阶段（即你负责的若干层）的前向计算。
-
通过PP组通信，将当前部分层的激活值传输给后继阶段的GPU
-
-
反向传播：
-
-
通过PP组通信，拿到后继层的梯度
-
通过TP组通信，完成当前部分层的反向传播
-
通过PP组通信，将当前部分层的梯度传输给前驱层
-
通过DP组通信，将所有DP组计算得到的梯度通过Reduce-Scatter进行全局平均，并按照预先约定好的DP分片规则， 每张卡只保留自己所需的全局平均梯度
-
-
参数更新：
-
-
**你现在有模型并行分割出的子模型的一部分模型优化器和对应的部分全局平均梯度，**更新这部分参数
-
通过DP组通信，通过All-Gather, 将DP组其他成员负责更新的模型参数同步过来，用于下一次前向传播
-
3. 使用ZeRO-2 DP(DP组的每一个成员保存一部分梯度和优化器状态)（不再兼容PP）
-
前向传播：
-
- 通过TP组通信，在TP组内协同完成整个模型的前向传播（TP组内每个GPU负责模型的一个切片）。
-
反向传播：
-
- 通过TP组通信，在TP组内协同完成整个模型的反向传播
-
-
- 在计算过程中（通常按层或块），一旦计算出某一部分局部梯度，立即通过DP组通信，将所有DP组计算得到的梯度通过Reduce-Scatter进行全局平均。操作结果和ZeRO-1相同，DP组的每一个成员只保留自己管理的这部分优化器状态对应的全局梯度
-
-
参数更新：
-
-
**你现在有模型并行分割出的子模型的一部分模型优化器和对应的一部分全局平均梯度，**更新这部分参数
-
通过DP组通信，通过All-Gather, 将DP组其他成员负责更新的模型参数同步过来，用于下一次前向传播
-
4. 使用ZeRO-3 DP(DP组的每一个成员保存一部分模型、梯度和优化器状态)
-
前向传播：
-
-
进行当前计算部分（如层）前，就通过DP组All-Gather通信，临时获取这部分计算对应的完整模型参数
-
通过TP组通信，在TP组内协同完成整个模型的前向传播
-
每个计算部分计算完成立刻丢弃模型参数，但保留激活值（张量并行分片后）
-
-
反向传播：
-
-
每开始一部分梯度的计算，就再次通过DP组All-Gather通信，临时获取这部分计算对应的完整模型参数
-
通过TP组通信，利用前向传播保存的激活值计算梯度，在TP组内协同完成整个模型的反向传播，计算完成立刻丢弃模型参数（张量并行分片后）
-
每完成一部分局部梯度的计算，立即通过DP组通信，将所有DP组计算得到的梯度通过Reduce-Scatter进行全局平均，DP组的每一个成员只保留自己管理的这部分优化器状态对应的全局梯度
-
-
参数更新：
-
-
你现在有模型并行分割出的子模型的一部分模型优化器和对应的部分全局平均梯度和对应的部分参数，直接更新即可
-
不再需要DP组通信
-
五、一些曾经困扰我的问题
1. ZeRO-3也能起到降低单卡显存压力的作用，也能将模型参数分片，为什么还需要TP？
ZeRO-3确实显著降低了模型参数、梯度和优化器状态占用的显存。然而，在训练真正超大规模模型时（如千亿、万亿参数），前向传播和反向传播中产生的中间计算结果——激活值(Activations)——所占用的显存会变得极其巨大，甚至超过模型参数本身。
-
ZeRO系列优化天然不会对激活值进行分片。在DP组内，每个成员都需要为相同的输入计算并存储一整套（完整且可能巨大的）激活值。
-
TP（以及PP）的关键优势在于，它们能将激活值本身也分割存储在TP组内的不同GPU上。每个GPU只存储激活值的分片，从而大幅降低单卡所需存储的激活值大小。
-
性能优势： TP可以利用单服务器内GPU间超高速互连（如NVLink）进行频繁的数据交换（同步模型层内切分计算产生的中间结果），效率通常远高于跨服务器通过网络进行的ZeRO通信（All-Gather）。实际测试（两块3060测试）表明，TP的吞吐量可达ZeRO-3 DP的1.5倍（大约1200 Toks/s vs 800 Toks/s）。
因此，TP（或PP）是解决大规模模型激活值显存瓶颈和提升单机计算效率的必要手段，ZeRO无法完全替代。
2. 为什么前几年deepspeed的老版本一直没有支持张量并行，与ZeRO系列优化相比他难在哪里。
DeepSpeed早期主打其ZeRO数据并行优化系列。其TP和PP支持常借助Megatron-LM等实现。直到最近的版本0.16.4（我在使用的是0.16.7）才支持原生的张量并行。主要原因在于：
-
模型侵入性不同：
-
-
ZeRO是一种数据并行优化策略，它对模型的结构本身不做任何修改。它的优化集中于管理模型状态（参数/梯度/优化器）在DP组内的分布和通信。对于GPU来说，DP组内其他卡的显存就像一个“远程缓存”，需要时通过通信取回参数。
-
TP和PP则是模型并行技术，它们直接改变了模型结构的计算图执行方式。TP需要重写模型层的实现，使其支持跨GPU的分割计算。PP需要将模型切割成段并在GPU间调度。
-
-
集成复杂性： 支持TP/PP需要框架能够自动（或通过用户注解）识别模型结构，并进行复杂的模型转换（图变换、层替换、切割、流水线编排等）。这比实现ZeRO复杂的通信逻辑更具挑战性。早期的DeepSpeed选择优先做好ZeRO，后来通过集成或自研支持了更完整的模型并行能力。
3. 为什么说ZeRO-1相对于ZeRO-0没有通信量增加，难道不是多了一次DP组间通信吗？
-
问题的核心在于：All-Reduce = Reduce-Scatter + All-Gather。
-
-
一个完整的All-Reduce操作（计算全局总和并分发到所有参与者）在通信量上等价于依次执行一个Reduce-Scatter（计算全局总和并分发部分结果给各参与者）和一个All-Gather（收集所有参与者的部分结果，组合成完整结果发给所有参与者）。
-
在带宽受限场景下（Bandwidth-bound），每个操作（Reduce-Scatter或All-Gather）的通信量大约是
M(这里M通常指模型参数总量 或者 梯度张量总量)。
-
-
具体分析：
-
-
ZeRO-0： 一次
All-Reduce(梯度平均) ，总通信量 ≈ 2M (每个GPU收发约 2 *M的数据量)。这是非ZeRO DP的标准做法。 -
ZeRO-1： 一次
Reduce-Scatter(梯度归约切分) [通信量 ≈ M] + 一次All-Gather(参数更新后同步) [通信量 ≈ M]。总通信量 ≈ 2M。虽然通信次数多了1次，但总字节数和ZeRO-0相同。只是分成了两步更小的通信。 -
ZeRO-2： 梯度
Reduce-Scatter[通信量 ≈ M] + 参数All-Gather[通信量 ≈ M]。总通信量 ≈ 2M。（梯度分片Reduce-Scatter在反向传播过程中提前被调用，更频繁触发小通信，但累积总量仍约为 M）。 -
ZeRO-3： 前向
All-Gather(参数) [通信量 ≈ M] + 反向All-Gather(参数) [通信量 ≈ M] + 梯度Reduce-Scatter[通信量 ≈ M]。总通信量 ≈ 3M。是ZeRO系列中通信量最大的方案。
-
-
核心结论： ZeRO-1相对于ZeRO-0，通信总字节数并未增加（都是2M），但通信模式发生了变化（从1个All-Reduce变成1个Reduce-Scatter + 1个All-Gather），通信次数增加了一次。ZeRO-3则在此基础上增加了50%的通信量。
4. 为什么ZeRO-2以上与PP并行不兼容
-
主要矛盾在于流水线微批次（Micro-batches）的执行策略与ZeRO的高频通信需求：
-
PP的执行特点： 为了最大化PP的GPU利用率（减少流水线”气泡“），实践中会将一个完整的”宏批次“（Macro-batch）拆分成
N个连续的微批次。PP组内的GPU并行处理不同的微批次（流水线填充态），并累积N个微批次计算产生的梯度，在一个宏批次结束时才对累积的梯度进行一次同步更新（类似于使用N倍大的batch size做了一次参数更新）。 -
ZeRO-2/3的通信要求： ZeRO-2需要在梯度计算过程中（可能每层）频繁地进行
Reduce-Scatter来切分/同步梯度。ZeRO-3更是需要在每次计算部分参数（如层）前都进行All-Gather来获取参数。 -
冲突： 在PP中立即应用ZeRO-2/3通信：
-
-
会打断流水线的连续执行（每次微批次结束都必须停下来等待DP跨组通信）。
-
如果要求每个微批次计算出的梯度都立即进行ZeRO通信（
Reduce-Scatter），那么总共需要的通信次数将是微批次数量（N）倍，导致通信量剧增（对于ZeRO-2是N * M梯度通信，对于ZeRO-3是2 * N * M甚至更多）。 -
这种通信开销通常巨大得不偿失，失去了PP高效利用计算资源的优势。
因此，在混合并行方案中，ZeRO-2/3通常只与TP结合使用，而避免与PP一同使用。
-
5. 现在大家通常采用的并行策略是什么
参考GPT-3、BLOOM、Qwen等模型的实践经验，一个高效通用的并行策略通常遵循以下原则：
-
TP在高速互联环境： 张量并行(TP)需要在层内进行极高频、低延迟的数据交换（如同步矩阵乘法的分块结果）。因此最适合部署在单个服务器节点内具有超高速互连(NVLink) 的多张GPU之间。例如，一台8卡服务器可能部署一个8卡的TP组来容纳一层。
-
PP跨节点扩展： 流水线并行(PP)阶段间的通信次数相对较少（主要是传输激活值和梯度），但每次传输的数据量很大。这种通信模式适合利用跨服务器的高速网络(如InfiniBand) 进行。通常，不同阶段分配到不同服务器节点。
-
DP扩展数据量： 当模型的参数量通过TP+PP的组合已经被充分拆分到多个GPU（服务器）上，使得单卡/服务器可以容纳模型的一个切分副本后。如果还有富余的服务器资源（包含这些TP+PP形成的组），就可以将每组TP+PP视为一个复制单元（副本），在其上启用数据并行(DP)，形成多个DP组。这允许使用更大的全局Batch Size进行训练。
-
在启用PP并行以后，则至多只能使用ZeRO-1 DP优化，否则会产生不可接受的通信开销放大。
六、总结
在大模型训练过程中，每块GPU虽然承担的是不同的计算任务，但是如果只从其“计算状态”和“通信状态”的角度考虑，其实每张卡做的事情都是完全一致的。这也就是为什么我们能够利用Deepspeed，Megatron-LM等编写一份代码，就能控制所有的GPU一起协同工作，训练出完整的大模型，而无需细节微操，为每一块GPU编写不同的代码。