在开始下一个阶段的开发之前，要先打好理论地基。虽然这部分之前零零散散的学过很多，但在通常的技术博客中通常只讲解每一种并行模式的技术原理和思想，将其和通信模式逐一详细联系起来并且做通信次数、数据量整合，甚至和反向传播的训练过程整合起来的并不多。

这篇文章打算稍微弥补一下这个问题，把所有零散的、常见或者不常见的并行切分方案，并且逐一评估其通信的需求、消耗、适用情况，在我的单人用户场景下，是否值得实现、实现难度如何。下面是总结。如果后面遇到了（甚至产生了）新的并行方案，也会一并记录在这里。

此外，在后面一篇博客当中，也会介绍一下每一种通信模式，以及通信模式对应的集群拓扑结构和相应常见的具体实现有哪些。

1、Tensor Parallelism（TP）— 层内切分

这个应该是最经典的并行方式了。 它的切分方式就是把每一层的权重矩阵沿着head或者hidden维度切到多卡上，每卡算一部分，合并结果。不同的卡持有不同的head、或者FFN的权重矩阵的不同部分。

特别值得一提的是，Megatron-LM中一种经典的做法是，FFN中有两个线性层，可以把第一个线性层（W_gate和W_up）按列切、第二个按行切。这样的好处是，在中间的计算过程中，中间结果不需要通信。

它的流程如下：先用被切分的，只有自己这部分head的W_Q、W_K、W_V矩阵乘以输入得到自己head的Q、K、V，然后用Attention算出自己head的输出，然后把自己的head送入按照同样模式切分的、自己本地的相应的W_o部分做乘法，得到一个部分和，然后进行all reduce。值得注意的是，由于切分是相同的，因此绕过了单卡情况下，对head的全局cat操作因此无需通信，只需要cat自己本地的head就可以了。all reduce之后，送入自己本地的第一个FFN权重，然后逐元素激活，送入第二个权重，算出本地部分和，然后进行第二次all reduce。

通信模式：All Reduce

通信次数：每个Transformer Block里有两次All-Reduce（Attention之后一次，FFN之后一次）。总共次数为2*模型层数。

通信量：每一层中，涉及到的需要通信的，相关的每层数据量 = batch_size × seq_len × hidden_size × dtype_bytes。在整个模型中，总数据量 = 每层数据量 × 2（每层两次） × layer_num（层数）。对于ring All-Reduce的情况，总通信量为总数据量 × TP 度 = 总数据量 × 2 × (n-1)/n。

涉及的模块：Attention、FFN

优点： 1、可以达到很细的切分粒度，因为不同Head、FFN权重都可以切分，每张卡的显存占用可以按 TP 度线性减少。 2、实现很成熟，是最早的并行模式之一，有现实的工业级别落地。 3、天然是完全负载均衡的。

缺点： 1、通信频率很高，对卡间带宽要求很高，产生的网络延迟在非NVlink的机内和机间基本不可接受 2、切分过细的矩阵可能导致GEMM的算术强度降低，从而降低效率

训练中反向传播的情况

Megatron-LM给出了一个很漂亮的结论：对于所有前向的All-Reduce情况，在反向中的算子为Identity（不作数学操作，仅把同一个输入x复制到每张卡上）；对于所有的前向的Identity情况，在反向中的算子为。因此，在反向传播中，通信产生的次数和前向传播一模一样、数据量也相同，但出现的位置和前向传播相反。前向传播在结尾处做all reduce，而反向传播在开始处（当然同时是反向传播自己的结尾处）做all reduce。

这其实很符合直觉，简单推导即可得到，在此按下不表。

需要注意的是，反向传播的一个stage，在一般论文里大约占据两个slot，也就是两个时间步。换言之，被建模为前向传播的时间的两倍。

2、Pipeline Parallelism（PP）— 层间切分

Layer 分离 — 切分不同的Transformer Block

最简单直接的并行模式。也就是直接把模型的不同层分配到不同卡上。例如，对于一个32层模型进行对4卡的PP，那么每一张卡恰好持有 8层。数据从第一卡计算完成后送入下一张卡，以此类推直到流向最后一卡，像流水线一样。

通信模式：点对点（Point-to-Point）

通信次数：整个前向传播过程中总通信次数N-1次，N为卡数。

通信量：每次切换卡，单次卡间通信量 = batch_size × seq_len × hidden_size × dtype_bytes，一次前向传播的总通信量 = num_gpus × 单次卡间通信量。

涉及的模块：不涉及具体模块，按卡的stage进行划分，粒度为layers。

优点： 1、通信量小、频率低，可以其通信量和频率要求允许跨节点进行通信。 2、实现简单、计算规整，可以保证层内计算和单卡完全相同，不产生任何浪费和算术强度降低。 3、反向传播计算流和单卡情况相同，不需要对反向传播的组织模式进行特别的处理。

缺点： 1、存在流水线空泡（Pipeline bubble）问题，流水线的启动。 2、对于训练场景有额外的问题：反向传播需要反向流水线，交叉进行在简单的调度策略下会造成更多空泡，需要micro-batch、1F1B调度等等更复杂的策略。 3、对于层数不能被卡数整除的场景，会产生负载不均衡。此外，首尾层的embedding/lm_head可能产生lagging。

训练中反向传播的情况

Pipeline Parallelism的反向传播在数学上和单卡完全一致，每个stage各自做标准的反向传播就行。但难点在于调度：怎样安排各stage的前向和反向的执行顺序，使得GPU尽量少空闲。

需要了解的调度策略如下：

1、GPipe 朴素方案。先把所有 micro-batch 的前向做完，再统一做反向。会产生很多空泡，而且需要保存大量的micro-batch的激活值。值得注意的是：重计算想法（只存储部分中间激活值以节省显存）也是在这部分研究当中一同产生的。

2、1F1B（One Forward One Backward） 在第一个forward做完之后立刻开始第一个backward的反向过程，和后面的forward过程穿插。通过组织调度数学上可以保证每个stage在任意时刻最多只需缓存PP个micro-batch的激活值。它的核心是：在稳定阶段，每个

3、Interleaved 1F1B 每个GPU不再持有连续的若干层而是持有总数相同但是在模型前向传播过程中不连续的层。这等价于增加了stage的数量，只不过可能有多个stage在同一张卡上。比如PP=4的情况下，8层模型、interleave=2；那么第一个stage（gpu）持有0、4层，第二个stage持有1、5层，第三个stage持有2、6层，第四个stage持有3、7层。 它的好处是可以通过很简单的方式增加流水线深度，从而允许更少的流水线空泡，坏处则很明显：大大增加了需要的通信量，进而增加了设备需求和延迟。

4、Zero-Bubble Pipeline 这是一个很巧妙的想法。它把反向传播的过程中的成分做了更细粒度的划分。具体来说，反向传播其实包含两个独立的计算：

• B（backward for input）：计算 ∂L/∂x，即对输入激活的梯度
• W（backward for weight）：计算 ∂L/∂W，即对权重的梯度

前者被前一个stage依赖的，前一个stage只有获得了B才能开始自己的反向传播计算过程。而后一个实际上是悬挂在依赖链上的分支，它依赖后一个stage的结果但本身并不被其他人依赖。W的结果只是累加到本stage的grad buffer里，等最后optimizer.step() 时用。它什么时候算都行，只要在optimizer.step()之前算完就行。那么一个自然的想法就是：能否先算有依赖的B，让后面的先开始，在有空的时候自己慢慢算W？答案是可以的。

通过手工调度，或者基于**ILP（整数线性规划）**求解的自动调度，可以最小化bubble的数量，甚至于接近0bubble的比例。当然，它实际上会引入一个问题：由于W延后，实际上一个step的总Latency变长了。但通过约束，可以完全保证跨step重叠至多在一个step以内，对训练收敛几乎没有可观测的影响。

Prefill-Decode 分离（PD Disaggregation）— 切分请求的不同阶段

Prefill 和 Decode 的计算特性完全不同：

Prefill:  大矩阵乘法，compute-bound，高算术强度，适合高算力卡
Decode:   矩阵-向量乘，memory-bound，低算术强度，适合高带宽卡

混在一起跑会互相干扰。Decode的延迟敏感请求被Prefill的大计算堵住，Prefill的高吞吐被Decode的频繁小请求打断。

更进一步的说，两者的基本的算术强度就不同：Decode的算术强度一般可以估算为，FFN部分为engine中的batch数量的两倍（fp16或bf16意义下。fp8则乘以2）、Attention部分固定为1。实测数据显示，总体强度一般是40-90 FLOPs/byte左右，远低于H100 BF16的ridge point 295。而Prefill的算术强度随着seq_len同步线性增长，一般远远超过任何gpu的ridge point。

Decode 阶段

线性层（FFN、QKV 投影、Output 投影）

Decode 每步只生成 1 个 token，所以是矩阵-向量乘法（GEMV）：

权重 W: [d, d]
输入 x: [batch_size, d]   （每个请求只有 1 个 token）

FLOPs:     2 × batch_size × d²
数据搬运:   d² × bytes_per_param（权重）+ batch_size × d × bytes（输入，通常可忽略）

权重远大于输入，所以：

算术强度 ≈ 2 × batch_size / bytes_per_param

Attention（decode 的 score 计算）

每个请求要用 1 个 query token 去和自己的整个 KV cache 做 attention：

Q: [batch_size, 1, d]
K cache: [batch_size, seq_len, d]    ← 每个请求有独立的 KV cache

FLOPs:     2 × batch_size × seq_len × d
数据搬运:   batch_size × seq_len × d × bytes（KV cache，每个请求各自的）

注意 KV cache 不能跨请求共享（每个请求的上下文不同），所以 batch 增大时数据搬运也等比增大：

算术强度 ≈ 2 / bytes_per_element ≈ 1（BF16）

不随 batch size 增长。

而Prefill的算术强度可以如此计算：

Prefill 阶段

线性层

Prefill一次处理整个 prompt，所以是矩阵-矩阵乘法（GEMM）：

权重 W: [d, d]
输入 X: [batch_size × seq_len, d]

FLOPs:     2 × batch_size × seq_len × d²
数据搬运:   d² × bytes_per_param（权重）
           + batch_size × seq_len × d × bytes（输入激活）

当 batch_size × seq_len 足够大时，权重搬运的成本被大量计算分摊掉：

算术强度 ≈ 2 × batch_size × seq_len / bytes_per_param

这些数字远远超过任何 GPU 的 ridge point。所以 Prefill 的线性层基本都是 compute-bound。

Attention（prefill 的 self-attention）

Prefill 中 Q、K、V 的长度都是 seq_len：

Q: [batch_size, seq_len, d]
K: [batch_size, seq_len, d]

Q×K^T 的 FLOPs:     2 × batch_size × seq_len² × d_head × num_heads
数据搬运:            batch_size × seq_len × d × bytes × 3（Q, K, V）

算术强度大约是：

算术强度 ≈ seq_len / (3 × bytes_per_element)

这些差异是如此的明显，以至于这个disaggreation是最先被提出的。因此，分离之后各自可以针对性地优化硬件配置和调度策略。

Attention-FFN分离（AF Disaggregation）— 切分Transformer的不同阶段

AF分离是一个相对来说更新的分离模式，而其和MoE的诞生关系密切，是一个很细粒度的分离。这是因为Expert的“专家”指的就是不同FFN的专家，它们共享Attention的所有参数，但各自具有各自的FFN层和相应参数。更进一步的说，两者占据显存的东西并不相同，Attention部分主要是因为KV Cache占据了大量显存，KV cache在不同卡之间的来回拷贝可能造成Attention性能的严重下降；而FFN中，权重占据了显存的绝大多数（甚至可以达到80%以上）。对于一个MoE来说，显存的膨胀同时也允许（或者说需求了）更细粒度的分离模式。此外，MoE的通信模式本身就需要大量的、不规则的、不可预知的点对点通信，Attention→FFN的传输可以和All-to-All的dispatch合并成一步，这意味着AF分离的边际代价相较于dense模型的情况减少了。反正都是无法避免要付出的，不如一口气全做完。

总的来说，Attention更倾向于小TP度（保持GEMM效率，减少KV cache副本和通信，而且Attention并没有非常庞大的参数量导致的需求），而FFN/Expert想要大EP度（容纳更多expert）。可以说MoE天然适合AF分离——既在动机上需要它（Attention和Expert的并行度不匹配），又在代价上容忍它（All-to-All吸收了额外通信）。

3、Data Parallelism（DP）— 请求间切分

这应该是最朴素的并行了，朴素到不需要怎么讲。它的复杂度不在于推理（推理就完全是不同的实例在独自工作，无通信），而是在于训练的冗余参数存储的优化，这里诞生了许多大模型训练重要的早期成果——Google Brain、ImageNet、GPT-2、BERT，一直到单个模型在单卡上彻底放不下为止。

通信模式：推理无（训练则是All Reduce（经典训练DP）/All-Gather（ZeRO-1）/Reduce-Scatter（ZeRO-2）/两者（ZeRO-3））

通信次数：推理0，训练则是每个训练步 2 次集合通信（1 次 Reduce-Scatter + 1 次 All-Gather）（ZeRO-1）/ L+1 次集合通信（L 次 Reduce-Scatter（反向传播中逐层做）+ 1 次 All-Gather（更新后同步权重））（ZeRO-2）/ 3L 次集合通信（前向 1 次 All-Gather + 反向 1 次 All-Gather + 反向 1 次 Reduce-Scatter）（ZeRO-3）。

通信量：推理0，训练则是2 × Φ（可训练参数量） × dtype_bytes（经典DP和ZeRO-1和ZeRO-2）/ 3 × Φ（可训练参数量） × dtype_bytes（ZeRO-3）。

涉及的模块：推理无，训练涉及到优化器状态（ZeRO-1）、梯度状态（ZeRO-2）和权重（ZeRO-3）。

优点： 1、无敌简单和容易理解。 2、推理完全无通信开销。 3、天然是完全负载均衡的。

缺点： 1、对大模型的一个卡装不下的情况毫无帮助。 2、训练的反向传播过程依然需要通信。

训练中反向传播的情况

这是ZeRO-1、2、3和FSDP大显身手的地方。

ZeRO-1：切分优化器状态

切分方式：优化器状态（Adam 的 m 和 v）按参数均分到 N 张卡。权重和梯度每卡完整保存。

通信模式：Reduce-Scatter + All-Gather

通信过程：

反向传播结束后:
  1. Reduce-Scatter 梯度
     每卡得到自己负责的那 1/N 参数的聚合梯度
     通信量: Φ × dtype_bytes（每卡发出完整梯度，收到 1/N）

  2. 每卡用聚合梯度更新自己那 1/N 的参数（本地操作，无通信）

  3. All-Gather 更新后的权重
     每卡把自己更新的 1/N 权重广播出去，收集到完整权重
     通信量: Φ × dtype_bytes

通信次数：每个训练步 2 次集合通信（1 次 Reduce-Scatter + 1 次 All-Gather）

每步总通信量：2Φ × dtype_bytes（和经典 DP 的 All-Reduce 相同）

ZeRO-2：切分优化器状态 + 梯度

切分方式：优化器状态和梯度都按参数均分到 N 张卡。权重每卡完整保存。

通信模式：Reduce-Scatter + All-Gather

通信过程：

反向传播中（逐层）:
  1. 每层算完梯度后立刻做 Reduce-Scatter
     每卡只保留自己负责的 1/N 梯度，其余丢弃
     → 梯度显存从 2Φ 降到 2Φ/N

反向传播全部结束后:
  2. 每卡更新自己负责的 1/N 参数（本地操作）

  3. All-Gather 更新后的权重
     通信量: Φ × dtype_bytes

通信次数：每个训练步 = L 次 Reduce-Scatter（反向传播中逐层做）+ 1 次 All-Gather（更新后同步权重）

每步总通信量：2Φ × dtype_bytes（仍然和经典 DP 相同）

ZeRO-3 / FSDP：切分优化器状态 + 梯度 + 权重

切分方式：优化器状态、梯度、权重全部按参数均分到 N 张卡。每卡只存 1/N。

通信模式：All-Gather + Reduce-Scatter（前向和反向每层都要做）

通信过程：

前向传播（逐层）:
  对于第 i 层:
    1. All-Gather 收集完整权重   通信量: Φ_i × dtype_bytes
    2. 用完整权重做前向计算
    3. 丢弃非本卡的权重（释放显存）

反向传播（逐层，从后向前）:
  对于第 i 层:
    4. All-Gather 再次收集完整权重  通信量: Φ_i × dtype_bytes
    5. 计算梯度
    6. 丢弃非本卡的权重
    7. Reduce-Scatter 梯度         通信量: Φ_i × dtype_bytes
       每卡只保留自己负责的 1/N

全部反向结束后:
  8. 每卡更新自己的 1/N 参数（本地操作，无通信）
  （不需要额外 All-Gather 同步权重，因为下一步的前向会自动收集）

通信次数：每层 3 次集合通信（前向 1 次 All-Gather + 反向 1 次 All-Gather + 反向 1 次 Reduce-Scatter），总计 3L 次

每步总通信量：

前向 All-Gather:    Φ × dtype_bytes   （所有层加起来）
反向 All-Gather:    Φ × dtype_bytes
反向 Reduce-Scatter: Φ × dtype_bytes
──────────────────────────────────────
总计:               3Φ × dtype_bytes

比经典 DP / ZeRO-1 / ZeRO-2 的 2Φ 多了 50%。多出来的就是前向传播中那一次 All-Gather——因为权重也切了，前向时也要收集。

4、Expert Parallelism（EP）— 专家参数间切分

这个同样是相对来说很容易理解的一种并行模式，也是MoE的必然选择，可以说这个并行模式和MoE架构是严格相互绑定的。这么做的主要原因有两个：第一，Expert的总量非常大，导致相同量级的MoE模型，计算量低得多，但所有Expert加起来的参数量远大于dense模型，导致模型根本不可能放在同一张卡上；第二，这些expert之间是独立的，不同token去不同的expert，中间过程不需要通信或者拼接，天然适合分布到不同卡上。

通信模式：All-to-All-v（数据量不规则）

通信次数：每个 MoE 层: 2 次 All-to-All（dispatch + combine），非 MoE 层（如 Attention）: 按 TP 的通信模式（All-Reduce）。

通信量：不可精确预测。每次All-to-All的通信量取决于token的分发模式，进一步地取决于router。理想均匀情况下:

每个 GPU 持有 batch_size × seq_len 个 token
每个 token 选 top-k 个 expert
总共需要分发的 token 数: batch_size × seq_len × top_k
其中留在本卡的比例: EP_degree 分之一
需要发送的: batch_size × seq_len × top_k × (1 - 1/EP) × hidden_size × dtype_bytes 

每个 MoE 层: 上述 × 2（dispatch + combine）

总通信量: 上述 × MoE层数。

但实际情况往往（高度）不均匀，热门 expert 收到的 token 多，冷门 expert 收到的少。

涉及的模块：FFN。

优点： 1、可以在不让计算量增加到难以承受程度的情况下，获得接近大参数量模型的性能。 2、切分粒度可以非常高。

缺点： 1、几乎总是一定有明显的负载不均衡，且无法在离线情况下静态的提前规划调度。 2、对显存的需求非常非常大。 3、训练比dense模型更加困难。

负载不均衡是MoE最难以处理的问题。Router是动态的，每个batch的token分发模式都不同。某些expert可能突然变"热门"，收到大量 token，而其他expert空闲。这会造成GPU的lagging情况，从而拖慢整个集群的性能。EP也是唯一一个通信模式动态变化的并行策略。这使得它的优化比其他并行困难得多，因为没法提前知道每一步的通信模式是什么样的。对它的调度策略比起前面的精密的线性规划模式，可能更接近于时序预测、控制论、排队论、资源调度等等范畴。可以想见，关于MoE模型的网络设计、集群设计、调度策略设计将会比前面的模式更有可以做的地方，可能也需要更长的时间才能在工业界收敛到一个相对来说足够优秀的统一策略，也可能根本不收敛。

训练中反向传播的情况

EP的训练问题并不在于技术上的问题——在DP、PP、TP等模式里已经基本上完全解决了。它面对的是专家坍缩和辅助负载均衡损失等等问题。它指的是一种Expert之间的马太效应：Router是可学习的，它在训练过程中会"偏心"——逐渐把越来越多的token路由给少数几个 expert，其他expert收到的token越来越少，最终几乎不参与计算。这是一个正反馈循环：某个expert因为偶然收到更多token，得到更多训练，变得更强，router就更倾向于选它，它就收到更多token……最终模型退化成一个dense模型，MoE的容量优势完全消失。更详细的介绍超出了这篇博客的范畴，因此就不多讲了。

此外，值得一提的是，Router的top-k选择是一个离散操作（不可微），但我们需要梯度回传。常用的做法是：

• 对被选中的expert，正常回传梯度
• 对没被选中的expert，梯度为零（它们根本没参与计算）
• Router的梯度通过加权系数（gating weight）回传——这些系数是连续的、可微的

这意味着router的训练信号只来自被选中的expert，它对未选中expert的"好坏"一无所知。这也是expert collapse的一个根源。

5、Sequence Parallelism （SP） — 序列维度切分

SP是基于前文提到的TP的一种改进（同时也必须在TP的基础上做，无法单独存在。此外，它的名字实在是容易引起误解），它处理的是非FFN和非Attention的、没有被TP切割的地方，如LayerNorm的激活值、Dropout的激活值、残差连接的激活值等等，当然也是一个训练专门使用的优化手段。这是因为，激活值存储这个问题在推理里根本不存在。SP优美的地方在于，它不增加通信量，而是通过拆解通信原语的方式无代价的获得收益——也许可以被称为一种通信的Kernel Fusion？

回顾TP用到的通信模式：每个Attention/FFN模块结束时做一次All-Reduce。而在底层通信原语的逻辑上，All-Reduce其实一般是拆成两步实现的：

All-Reduce = Reduce-Scatter + All-Gather

SP 把这个拆分利用起来了：

不用 SP 时:
  TP 模块结束 → All-Reduce → 每卡得到完整结果 → LayerNorm(完整)

用 SP 时:
  TP 模块结束 → Reduce-Scatter → 每卡得到 1/N 的结果 → LayerNorm(1/N)
  → 下一个 TP 模块开始前 → All-Gather → 收集完整输入 → TP 计算

通信模式：Reduce-Scatter + All-Gather（替代 TP 原本的 All-Reduce）

通信次数：和 TP 相同，每个 Transformer Block 4 次集合通信（前向 2 次 + 反向 2 次），只是把 All-Reduce 拆成了 RS + AG

总通信量：Reduce-Scatter + All-Gather = All-Reduce。通过把All-Reduce拆分和还原，做到了在不增加额外开销的情况下减少大量激活值存储。

6、Context Parallelism （CP） — 长序列 Attention 切分

这是最近比较新的一种并行模式。CP应该是目前主流并行策略中最新加入的一维。Llama 3的训练就用了4D并行：TP + CP + PP + DP，CP 是第四维。

CP和Ring Attention技术密切相关。它是把输入序列沿token维度切分到多张卡上。每卡只持有序列的一段，以此实现利用并行性来加速的效果。

对于非Attention的操作（FFN、LayerNorm等），token之间没有交互，每卡独立计算自己的部分即可，完全不需要通信。唯一需要通信的是 Attention，因为每个token的Query需要和所有token的Key、Value交互。

Meta 在Attention模块中实现了两种CP变体，通常被称为Ring Attention。

Pass-KV（经典Ring Attention）

每卡保留自己的 Q 不动
K、V 块在 GPU 之间沿环传递
每轮：用本地 Q 和收到的 KV 块计算部分 Attention

标准Ring Attention，适合大多数场景。

Pass-Q（反向传递）

每卡保留自己的 K、V 不动
Q 块在 GPU 之间沿环传递
每轮：用收到的 Q 和本地 KV 计算部分 Attention

Pass-Q在某些特定场景下更优。当KV cache命中率低于5%时，Pass-KV更优。