大模型中的各种并行：TP DP EP PP
TP（Tensor Parallel）并行介绍
1. 为什么大模型需要 TP
大模型的 Transformer 层中，最主要的参数和计算量集中在矩阵乘法上，例如 Attention 中的 Q/K/V Projection、Output Projection，以及 MLP 中的两层 Linear。随着模型隐藏维度、注意力头数和 FFN 中间维度增大，单张 GPU 往往无法同时承担完整权重、激活和 KV Cache 的显存压力。
Tensor Parallel（TP，张量并行）就是把一个大矩阵、一个大张量或一组 Attention Heads 切分到多张 GPU 上，让多张 GPU 共同完成同一层的计算。它和 Pipeline Parallel（PP）不同：PP 是按层切分模型；TP 是在同一层内部切分矩阵计算。
在大模型推理和训练中，TP 常用于：
- 降低单卡权重显存占用。
- 降低单卡矩阵乘法计算压力。
- 支持更大的 hidden size、更多 attention heads、更大的 FFN intermediate size。
- 在单节点多卡或高速互联多卡场景下提升吞吐。
常见配置是 TP=2、TP=4、TP=8。例如一个 64-head 的模型，在 TP=4 时，每张 GPU 负责 16 个 attention heads。
2. TP 的基本思想
TP 的核心问题是：一个 Linear 层 Y = XW，如何把 W 切到多张 GPU 上，同时尽量减少通信。
假设：
- 输入
X的形状是[B, H]。 - 权重
W的形状是[H, O]。 - 输出
Y的形状是[B, O]。 - TP world size 为
N。
TP 通常使用两类切分方式：
- Column Parallel：按输出维度切分权重列。
- Row Parallel：按输入维度切分权重行。
3. Column Parallel Linear
Column Parallel 是把权重矩阵按列切分：
W = [ W0 | W1 | W2 | W3 ]
X: [B, H]
Wi: [H, O/4]
Yi: [B, O/4]
每张 GPU 都拿到完整输入 X，但只保存自己那一部分权重 Wi，计算自己那一部分输出 Yi。
X [B,H]
│
┌─────────┼─────────┬─────────┬
│ │ │ │
GPU0 GPU1 GPU2 GPU3
W0 W1 W2 W3
[H,O/4] [H,O/4] [H,O/4] [H,O/4]
│ │ │ │
Y0 Y1 Y2 Y3
[B,O/4] [B,O/4] [B,O/4] [B,O/4]
Column Parallel 的特点：
- 每卡只保存
1/N的输出列权重。 - 每卡只计算
1/N的输出通道。 - 如果下一层可以继续使用分片输出，则不需要立即 AllGather。
- 常用于 QKV Projection 和 MLP Linear1。
4. Row Parallel Linear
Row Parallel 是把权重矩阵按行切分：
┌ W0 ┐
├ W1 ┤
W = ├ W2 ┤
└ W3 ┘
Xi: [B, H/4]
Wi: [H/4, O]
Partial Yi: [B, O]
输入 X 也需要按 hidden 维度切分，每张 GPU 计算一个 partial output。因为完整输出是所有 partial output 的和，所以最后需要一次 AllReduce(SUM)。
X0 [B,H/4] ── W0 [H/4,O] ── Partial Y0 [B,O]
X1 [B,H/4] ── W1 [H/4,O] ── Partial Y1 [B,O]
X2 [B,H/4] ── W2 [H/4,O] ── Partial Y2 [B,O]
X3 [B,H/4] ── W3 [H/4,O] ── Partial Y3 [B,O]
│
AllReduce(SUM)
│
Y [B,O]
Row Parallel 的特点：
- 每卡只保存
1/N的输入行权重。 - 每卡输出 shape 仍是
[B, O]，但只是 partial result。 - 必须通过 AllReduce 把 partial results 求和。
- 常用于 Attention Output Projection 和 MLP Linear2。
5. Transformer Decoder Layer 中的 TP=4 示例
下面用一个简化 Decoder Layer 展示 TP=4 的数据流。假设：
- Hidden size = 8192。
- Attention 总 head 数 = 64。
- TP = 4。
- 每张 GPU 负责 16 个 Attention Heads。
- MLP intermediate size 被切到每卡 7168（总中间维度约为 28672）。
Input Hidden
Shape: [B, 8192]
│
│ Broadcast / Replicated Input
│ 每张 GPU 都看到同一份输入 hidden states
│
├──────────────┬──────────────┬──────────────┐
│ │ │ │
GPU0 GPU1 GPU2 GPU3
16 Heads 16 Heads 16 Heads 16 Heads
│ │ │ │
└────── QKV Projection + Attention ──────┘
│
5.1 QKV Projection：Column Parallel / Head Parallel
QKV Projection 通常可以按 attention head 切分。每张 GPU 只计算自己负责 heads 的 Q/K/V。
在这个例子中，模型一共有 64 个 attention heads，TP=4 表示把同一层的 attention heads 分到 4 张 GPU 上：
Total heads = 64
TP size = 4
Heads/GPU = 64 / 4 = 16
因此每张 GPU 负责 16 个 heads：
Input Hidden: [B, 8192]
GPU0: heads 0 - 15 -> 计算这些 heads 的 Q/K/V
GPU1: heads 16 - 31 -> 计算这些 heads 的 Q/K/V
GPU2: heads 32 - 47 -> 计算这些 heads 的 Q/K/V
GPU3: heads 48 - 63 -> 计算这些 heads 的 Q/K/V
如果每个 head 的 head_dim 是 128，那么每张 GPU 对应的 attention hidden 宽度是：
16 heads/GPU × 128 head_dim = 2048 hidden dims/GPU
8192 hidden dims / 4 GPUs = 2048 hidden dims/GPU
所以经过 QKV Projection 和 Attention 之后，每张 GPU 持有自己那 16 个 heads 的 attention 输出分片：
GPU0 attention output: [B, 2048]
GPU1 attention output: [B, 2048]
GPU2 attention output: [B, 2048]
GPU3 attention output: [B, 2048]
这一步通常不需要跨卡通信，原因是一个 attention head 的计算只依赖本 head 的 Q、K、V：
Q_i, K_i, V_i
│
Q_i × K_i^T
│
Softmax
│
Attention_i = Softmax(Q_i K_i^T) × V_i
也就是说，GPU0 计算 heads 0-15 时，不需要 GPU1/2/3 上 heads 16-63 的 Q/K/V。每个 GPU 可以独立完成自己负责 heads 的 QK^T、Softmax 和 AV 计算。
跨卡通信发生在后面的 Output Projection。因为 attention 最终要回到完整 hidden size [B, 8192]，而每张 GPU 当前只有 [B, 2048] 的 head 分片，所以 Output Projection 使用 Row Parallel，每张 GPU 先计算 partial output [B,8192]，再通过 AllReduce(SUM) 合成完整输出。
5.2 Attention Output Projection：Row Parallel
Attention 之后，各卡持有不同 heads 对应的 hidden 分片。Output Projection 要把这些分片映射回完整 hidden size。
Hidden 切分:
GPU0: [B, 2048]
GPU1: [B, 2048]
GPU2: [B, 2048]
GPU3: [B, 2048]
Output Projection 权重切分:
GPU0: [2048, 8192]
GPU1: [2048, 8192]
GPU2: [2048, 8192]
GPU3: [2048, 8192]
每张 GPU 计算一个 partial output：
GPU0 partial output: [B, 8192]
GPU1 partial output: [B, 8192]
GPU2 partial output: [B, 8192]
GPU3 partial output: [B, 8192]
│
AllReduce(SUM)
│
Full Hidden: [B, 8192]
这一步需要 AllReduce，因为完整输出等于四个 partial outputs 的逐元素求和。
5.3 MLP Linear1：Column Parallel
MLP 第一层通常从 hidden size 映射到 intermediate size。以 SwiGLU / SiLU Gate 结构为例，Linear1 可能包含 up projection 和 gate projection。
Input Hidden: [B, 8192]
MLP Linear1 权重切分:
GPU0: [8192, 7168]
GPU1: [8192, 7168]
GPU2: [8192, 7168]
GPU3: [8192, 7168]
每卡输出:
GPUi: [B, 7168]
这一步是 Column Parallel，每张 GPU 计算 intermediate hidden 的一部分。随后激活函数可以在本卡独立执行：
GPUi: [B, 7168]
│
SiLU / Gate
│
GPUi: [B, 7168]
因为 MLP Linear2 正好可以接收分片 intermediate hidden，所以这里通常不需要 AllGather。
5.4 MLP Linear2：Row Parallel
MLP 第二层把 intermediate hidden 映射回 hidden size。由于 intermediate hidden 已经按 GPU 分片，Linear2 适合使用 Row Parallel。
MLP Linear2 权重切分:
GPU0: [7168, 8192]
GPU1: [7168, 8192]
GPU2: [7168, 8192]
GPU3: [7168, 8192]
每卡 partial output:
GPUi: [B, 8192]
最后再做一次 AllReduce：
GPU0 partial [B,8192]
GPU1 partial [B,8192]
GPU2 partial [B,8192]
GPU3 partial [B,8192]
│
AllReduce(SUM)
│
Output Hidden [B,8192]
6. Decoder Layer TP=4 总流程图
Input Hidden [B,8192]
│
│ Broadcast / Replicated
│
┌──────┴────────┬──────────────┬──────────────┐
│ │ │ │
GPU0 GPU1 GPU2 GPU3
16 Heads 16 Heads 16 Heads 16 Heads
│ │ │ │
QKV + Attn QKV + Attn QKV + Attn QKV + Attn
│ │ │ │
[B,2048] [B,2048] [B,2048] [B,2048]
│ │ │ │
Output Projection Row Parallel
│ │ │ │
Partial [B,8192] Partial Partial Partial
└──────┬────────┴──────┬───────┴──────┬───────┘
│ │ │
└────── AllReduce(SUM) ────────┘
│
Full Hidden [B,8192]
│
MLP Linear1 Column Parallel
│
┌───────────────┼──────────────┬──────────────┐
│ │ │ │
[B,7168] [B,7168] [B,7168] [B,7168]
│ │ │ │
SiLU/Gate SiLU/Gate SiLU/Gate SiLU/Gate
│ │ │ │
MLP Linear2 Row Parallel
│ │ │ │
Partial [B,8192] Partial [B,8192] Partial [B,8192] Partial [B,8192]
└───────────────┴────── AllReduce(SUM) ──────┘
│
Output Hidden [B,8192]
6.1 按 Attention Head 展开后的 Shape 流程图
上面的流程图使用 [B,8192] 表示 hidden states，这是 Linear 层最常见的二维视角。对于 Attention 来说，也可以把同一个 hidden 维度按 head 展开：
hidden_size = head_nums × head_dim
8192 = 64 × 128
[B, 8192] <=> [B, 64, 128]
这里：
B表示 batch size，也可以理解为当前参与计算的 token 数，实际实现中常见形状还会包含 sequence 维度，例如[batch, seq_len, hidden_size]。64表示总 attention head 数。128表示每个 head 的维度head_dim。8192就是 input hidden size，等于64 × 128，所以 hidden size 并没有消失，只是被 reshape/view 成了 head 维度形式。
按 head 展开后，TP=4 的 Attention 部分可以这样理解：
Input Hidden
[B, 8192]
│
│ reshape / view for attention
▼
[B, 64, 128]
│
│ TP=4，按 head_nums 维度切分
│ 每张 GPU 拿 16 个 heads
│
├─────────────────┬─────────────────┬─────────────────┐
│ │ │ │
GPU0 GPU1 GPU2 GPU3
[B,16,128] [B,16,128] [B,16,128] [B,16,128]
heads 0-15 heads 16-31 heads 32-47 heads 48-63
│ │ │ │
Q/K/V Projection Q/K/V Projection Q/K/V Projection Q/K/V Projection
│ │ │ │
Q,K,V: Q,K,V: Q,K,V: Q,K,V:
[B,16,128] [B,16,128] [B,16,128] [B,16,128]
│ │ │ │
Attention Attention Attention Attention
(QK^T + Softmax + AV, each GPU local)
│ │ │ │
Attn output Attn output Attn output Attn output
[B,16,128] [B,16,128] [B,16,128] [B,16,128]
│ │ │ │
│ flatten local heads: [B,16,128] -> [B,2048]
▼ ▼ ▼ ▼
[B,2048] [B,2048] [B,2048] [B,2048]
│ │ │ │
Output Projection Row Parallel
│ │ │ │
Partial [B,8192] Partial [B,8192] Partial [B,8192] Partial [B,8192]
└─────────────────┴────── AllReduce(SUM) ────────┘
│
▼
Full Hidden [B,8192]
为什么图里可以写 [B,64,128]，但前面又写 [B,8192]？原因是它们表达的是同一份数据的不同视角：
Linear 视角: [B, hidden_size] = [B, 8192]
Attention 视角: [B, head_nums, head_dim] = [B, 64, 128]
在进入 Q/K/V 和 Attention head 计算时，通常会把 hidden size 拆成 head_nums × head_dim。TP 按 head 并行时，切分的是 head_nums 这一维，而不是把 head_dim=128 再切碎。因此 TP=4 时，每张 GPU 得到：
[B, 64, 128] / 4 GPUs = 每卡 [B, 16, 128]
每张 GPU 本地完成 16 个 heads 的 Attention 后，再把本卡 16 个 heads flatten 回 [B,2048]，接到后面的 Output Projection。
这里的 Output Projection 为什么可以使用 Row Parallel？关键是 Attention 输出已经天然按 hidden 输入维度分片了：
完整 Attention 输出如果合在一起，本来是：
X_attn: [B, 8192]
但 TP=4 后，每张 GPU 只持有其中 1/4：
GPU0: X0 [B, 2048]
GPU1: X1 [B, 2048]
GPU2: X2 [B, 2048]
GPU3: X3 [B, 2048]
Output Projection 的完整矩阵可以写成：
Y = X_attn × W_o
X_attn: [B, 8192]
W_o: [8192, 8192]
Y: [B, 8192]
因为 X_attn 已经按输入 hidden 维度切成了 4 份，所以 W_o 也沿输入维度，也就是矩阵的 row 方向切成 4 份：
┌ W_o0 ┐ W_o0: [2048, 8192]
├ W_o1 ┤ W_o1: [2048, 8192]
W_o [8192,8192] = ├ W_o2 ┤ W_o2: [2048, 8192]
└ W_o3 ┘ W_o3: [2048, 8192]
于是每张 GPU 只需要用自己的 attention 分片乘自己的 row-shard 权重：
GPU0: X0 [B,2048] × W_o0 [2048,8192] = Partial Y0 [B,8192]
GPU1: X1 [B,2048] × W_o1 [2048,8192] = Partial Y1 [B,8192]
GPU2: X2 [B,2048] × W_o2 [2048,8192] = Partial Y2 [B,8192]
GPU3: X3 [B,2048] × W_o3 [2048,8192] = Partial Y3 [B,8192]
完整矩阵乘法可以拆成四个 partial results 的和：
Y = X_attn × W_o
= X0 × W_o0 + X1 × W_o1 + X2 × W_o2 + X3 × W_o3
= Partial Y0 + Partial Y1 + Partial Y2 + Partial Y3
所以 Output Projection 这里使用 Row Parallel 是自然的：每张 GPU 的输入已经是 [B,2048] 分片，对应权重也切成 [2048,8192]，每卡算出 [B,8192] partial output，最后通过 AllReduce(SUM) 合成完整 [B,8192]。
7. 通信总结
| 阶段 | 常见切分方式 | 每卡持有内容 | 是否需要通信 |
|---|---|---|---|
| QKV Projection | Column Parallel / Head Split | 部分 QKV heads | 通常不需要 |
| Attention 计算 | Head Parallel | 部分 heads 的 attention | 通常不需要 |
| Output Projection | Row Parallel | 部分输入行权重 | 需要 AllReduce |
| MLP Linear1 | Column Parallel | 部分 intermediate 输出列 | 通常不需要立即 Gather |
| Activation / Gate | Local | 本卡 intermediate 分片 | 不需要 |
| MLP Linear2 | Row Parallel | 部分 intermediate 输入行 | 需要 AllReduce |
| LayerNorm / Residual | Replicated Hidden | 完整 hidden | 通常依赖前序 AllReduce 后本地执行 |
一个标准 Transformer Decoder Layer 中，TP 通常会产生两次主要 AllReduce：
- Attention Output Projection 之后。
- MLP Linear2 之后。
这也是 TP 性能优化的关键：矩阵乘法被分摊了，但通信开销会增加。
8. TP 的优点
TP 的主要优点包括：
- 单卡权重显存下降到约
1/TP。 - 单卡 GEMM 计算量下降到约
1/TP。 - Attention heads 可以自然按 head 切分。
- 适合单节点多 GPU，尤其是 NVLink、PCIe Switch、IB/RDMA 等高速互联环境。
- 对大 hidden size、大 FFN intermediate size 的模型非常有效。
对于推理场景，TP 可以让单个大模型跨多卡部署。例如单卡放不下 70B 模型时，可以用 TP=4 或 TP=8 分摊权重。
9. TP 的代价和限制
TP 不是免费的，它会带来明显通信成本：
- Row Parallel 后需要 AllReduce。
- batch size 或 sequence length 较小时，通信延迟可能成为瓶颈。
- TP 跨节点时，如果网络带宽不足，性能可能显著下降。
- TP size 不能随意增大，通常要满足 attention heads、hidden size、intermediate size 能被 TP 整除。
- 每层都涉及多卡协同，调试和 profiling 难度高于单卡。
实践中，TP 最适合放在高速互联的同一节点内。例如 8 卡服务器常用 TP=8 或 TP=4。跨节点扩展时，通常会组合 PP、DP、EP 等其他并行方式。
10. 大模型中的其他常见并行方式概括
前面重点介绍了 TP。真实大模型训练和推理系统通常不会只使用 TP，而是把 TP、DP、PP、EP、SP 等并行方式组合起来。不同并行方式切分的对象不同，解决的问题也不同。
| 并行方式 | 全称 | 主要切分对象 | 主要解决的问题 | 典型通信 |
|---|---|---|---|---|
| TP | Tensor Parallel | 单层内部的矩阵、hidden 维度、attention heads | 单层太大，单卡放不下或算不动 | AllReduce / AllGather / ReduceScatter |
| DP | Data Parallel | batch 数据 | 多副本训练，提高训练吞吐 | 梯度 AllReduce |
| PP | Pipeline Parallel | Transformer 层 | 模型层数太多，无法放在单卡或单节点 | stage 间 activation / gradient 传递 |
| EP | Expert Parallel | MoE experts | MoE expert 数量多，单卡放不下全部 experts | token dispatch / combine，AllToAll |
| SP | Sequence Parallel | sequence 维度上的激活 | 降低长序列训练中的激活显存 | AllGather / ReduceScatter |
可以用一句话理解它们的区别：
TP: 把一层内部切开。
DP: 把不同 batch 样本切开。
PP: 把不同 Transformer 层切开。
EP: 把 MoE 的不同 experts 切开。
SP: 把 sequence 维度上的激活切开。
一个真实的大模型系统可能采用类似组合：
节点内: TP=4 或 TP=8
节点间: DP 或 PP
MoE 层: EP
长序列训练: SP
全局 GPU 数 = TP × PP × DP × EP 等并行维度的组合
下面分别详细介绍 DP、PP、EP、SP。TP 已在前文详细展开，这里只作为对比。
11. DP：Data Parallel，数据并行
11.1 DP 解决什么问题
Data Parallel 是最经典、最容易理解的并行方式。它不切模型结构，而是复制多份完整模型，每张 GPU 或每组 GPU 处理不同的数据 batch。
假设有 4 张 GPU，DP=4：
完整模型 replica 0 放在 GPU0，处理 batch 0
完整模型 replica 1 放在 GPU1，处理 batch 1
完整模型 replica 2 放在 GPU2，处理 batch 2
完整模型 replica 3 放在 GPU3，处理 batch 3
每张 GPU 上的模型参数初始相同，但输入数据不同。前向和反向各自本地计算，反向结束后需要把各卡梯度同步，使所有 replica 的参数更新保持一致。
11.2 DP 的计算流程
Global Batch
│
│ split by batch dimension
│
┌──┴─────────┬───────────┬───────────┐
│ │ │ │
GPU0 GPU1 GPU2 GPU3
Batch 0 Batch 1 Batch 2 Batch 3
Model copy Model copy Model copy Model copy
│ │ │ │
Forward Forward Forward Forward
Backward Backward Backward Backward
│ │ │ │
Grad0 Grad1 Grad2 Grad3
└────── Gradient AllReduce / ReduceScatter ──────┘
│
Same averaged gradients
│
Optimizer step on each replica
11.3 DP 的通信
DP 主要通信发生在反向传播后：
各 GPU 本地梯度 -> Gradient AllReduce -> 所有 GPU 得到相同平均梯度
训练中通信量通常和参数量相关。模型越大，梯度同步成本越高。
推理中一般不需要梯度同步，所以“推理 DP”通常就是多份模型副本处理不同请求，更多是服务层面的 replica 并行。
11.4 DP 的优点和限制
优点：
- 概念简单，实现成熟。
- 适合提升训练吞吐。
- 扩展 batch size 很自然。
- 对模型结构侵入较小。
限制：
- 每张 GPU 都要保存完整模型参数，不能解决“单卡放不下模型”的问题。
- 训练时需要同步梯度，模型越大通信越重。
- 全局 batch size 变大后，可能需要调整学习率、warmup 和优化器参数。
DP 常和 TP、PP 组合使用。例如每个 DP replica 内部再用 TP=4，把一个大模型切到 4 张 GPU 上；然后复制多个这样的 TP group 做 DP。
DP=2, TP=4，总共 8 张 GPU
DP replica 0: GPU0-3 组成 TP=4
DP replica 1: GPU4-7 组成 TP=4
12. PP：Pipeline Parallel，流水线并行
12.1 PP 解决什么问题
Pipeline Parallel 是按 Transformer 层切分模型。它适合模型层数很多、完整模型无法放在一张 GPU 或一个 TP group 内的场景。
例如一个 80 层 Decoder-only 模型，PP=4 时可以这样切：
GPU group 0: layers 0 - 19
GPU group 1: layers 20 - 39
GPU group 2: layers 40 - 59
GPU group 3: layers 60 - 79
每个 pipeline stage 只保存自己负责的层。输入 hidden states 从 stage 0 开始，逐 stage 向后传递。
12.2 PP 的前向流程
Input tokens
│
Stage 0: Embedding + Layer 0-19
│ activation
▼
Stage 1: Layer 20-39
│ activation
▼
Stage 2: Layer 40-59
│ activation
▼
Stage 3: Layer 60-79 + LM Head
│
Logits
训练时反向传播沿相反方向传回 gradients：
Stage 3 -> Stage 2 -> Stage 1 -> Stage 0
12.3 为什么 PP 需要 micro-batch
如果只送一个大 batch，stage 0 算完后 stage 1 才开始，stage 2 和 stage 3 在前面阶段工作时会空闲，GPU 利用率低。为减少空闲，PP 会把 batch 切成多个 micro-batches，让不同 stage 同时处理不同 micro-batch。
时间片 1: Stage0 处理 micro-batch 0
时间片 2: Stage0 处理 micro-batch 1, Stage1 处理 micro-batch 0
时间片 3: Stage0 处理 micro-batch 2, Stage1 处理 micro-batch 1, Stage2 处理 micro-batch 0
...
这种方式像流水线一样让多个 stage 同时工作。
12.4 PP 的通信
PP 的通信主要是相邻 stage 之间传递 activation 和 gradient：
Forward: stage i activation -> stage i+1
Backward: stage i+1 gradient -> stage i
和 TP 的 AllReduce 不同，PP 通常是点对点通信。通信量和 activation size、micro-batch size、sequence length、hidden size 有关。
12.5 PP 的优点和限制
优点：
- 可以把很多层分散到多张 GPU 或多个节点。
- 能显著降低单设备模型参数显存。
- 适合超大 dense 模型训练。
限制：
- 有 pipeline bubble，stage 数越多、micro-batch 越少，空闲比例越高。
- 需要处理 stage 间负载均衡，不同层计算量不一致时会拖慢流水线。
- 推理 decode 阶段 token-by-token 生成时，PP 的流水线效率可能不如训练 prefill 阶段明显。
PP 常和 TP 组合：每个 pipeline stage 内部是一组 TP GPU，stage 之间再串成 pipeline。
TP=4, PP=2，总共 8 张 GPU
Stage 0: GPU0-3，负责前半层，内部 TP=4
Stage 1: GPU4-7，负责后半层，内部 TP=4
13. EP：Expert Parallel，专家并行
13.1 EP 解决什么问题
Expert Parallel 主要用于 MoE（Mixture of Experts）模型。MoE 层中有多个 experts，每个 token 只会被路由到其中少数几个 experts，例如 top-1 或 top-2。
Dense MLP 是所有 token 都经过同一组 MLP 权重；MoE MLP 则是有多个 expert MLP：
Dense MLP:
所有 token -> 同一个 MLP
MoE MLP:
token -> router -> expert 0 / expert 1 / ... / expert N
当 expert 数量很多时，单卡保存全部 experts 会很浪费或放不下。EP 把 experts 分布到多张 GPU 上，每张 GPU 只保存一部分 experts。
13.2 EP 的基本流程
假设有 8 个 experts，EP=4：
GPU0: expert 0, expert 1
GPU1: expert 2, expert 3
GPU2: expert 4, expert 5
GPU3: expert 6, expert 7
MoE 层前向流程：
Input tokens
│
Router / Gate
│
为每个 token 选择 top-k experts
│
Token Dispatch
│
把 token 发送到对应 expert 所在 GPU
│
Experts 本地计算
│
Token Combine
│
把 expert 输出发送回原 token 位置并加权合并
13.3 EP 的通信
EP 的主要通信是 token dispatch 和 token combine，常用 AllToAll 类型通信。
Dispatch: token 从原 GPU 发送到 expert 所在 GPU
Combine: expert 输出再发送回原 GPU 或原 token 顺序
如果 router 选择 top-2 experts，一个 token 可能会被发送给两个 experts，最后再把两个 expert 输出按 gate weight 加权合并。
13.4 EP 的优点和限制
优点：
- 可以支持非常多 experts。
- 每张 GPU 只保存部分 experts，降低 expert 参数显存。
- MoE 激活参数量小于总参数量，每个 token 只计算少数 experts，计算效率高。
限制：
- token 路由会产生 AllToAll 通信，通信和负载均衡很关键。
- 如果 router 把大量 token 分到少数 experts，会出现 expert overload。
- 需要 capacity factor、token dropping、expert load balancing loss 等机制控制负载。
- 实现复杂度明显高于 dense MLP。
13.5 EP 和 TP 的关系
EP 切的是 experts，TP 切的是 expert 内部的矩阵。两者可以组合：
EP: 不同 GPU group 放不同 experts
TP: 每个 expert 内部的大矩阵再切到多张 GPU 上
例如一个 MoE 层有 64 个 experts，可以先用 EP 把 experts 分到不同 GPU，再对每个 expert 的 MLP 使用 TP。
14. SP：Sequence Parallel，序列并行
14.1 SP 解决什么问题
Sequence Parallel 主要用于降低训练时 activation 显存。TP 可以切权重和 hidden 维度，但某些操作仍然需要保存大量按 sequence 展开的激活，例如 LayerNorm、Dropout、Residual、部分 MLP 输入输出等。
当 sequence length 很长时，activation 显存可能非常高：
Activation size ≈ batch × seq_len × hidden_size
SP 的思路是把 sequence 维度也切到多张 GPU 上，让每张 GPU 只保存一部分 token 的激活。
14.2 SP 的 shape 示例
假设：
Activation: [B, S, H]
SP=4
按 sequence 维度切分：
GPU0: [B, S/4, H]
GPU1: [B, S/4, H]
GPU2: [B, S/4, H]
GPU3: [B, S/4, H]
这样每张 GPU 保存的 activation 显存约下降到 1/SP。
14.3 SP 的通信
SP 常和 TP 配合。某些操作可以在 sequence 分片上本地完成；某些操作需要完整 hidden 或完整 sequence，则需要 AllGather 或 ReduceScatter。
典型模式：
ReduceScatter: 把完整 activation 切成 sequence 分片
Local compute: 每卡处理自己的 sequence shard
AllGather: 在需要完整 activation 的地方重新收集
14.4 SP 的优点和限制
优点：
- 对长序列训练特别有价值。
- 降低 activation 显存，而不只是降低参数显存。
- 常与 TP 组合，进一步降低训练显存压力。
限制：
- 会引入额外 AllGather / ReduceScatter 通信。
- 对实现要求较高，需要明确哪些算子支持 sequence shard。
- 推理场景中是否需要 SP，取决于 batch、sequence length、KV Cache 和框架实现。
15. 多种并行方式如何组合
大模型并行通常是分层组合，而不是互相替代。一个常见理解方式是：
DP group 之间: 复制多份模型，处理不同数据
PP stage 之间: 不同层放在不同 stage
TP group 内部: 每一层的大矩阵切到多张 GPU
EP group 内部: MoE experts 分布到不同 GPU
SP: 在 TP 相关区域进一步切 sequence activation
例如 64 张 GPU 可以组合为：
TP=4, PP=4, DP=4
总 GPU = 4 × 4 × 4 = 64
每个 DP replica 有 16 张 GPU
每个 replica 内分 4 个 pipeline stages
每个 stage 内部用 4 张 GPU 做 tensor parallel
如果模型是 MoE，还可以加入 EP：
TP=4, PP=2, EP=4, DP=2
总并行规模 = 4 × 2 × 4 × 2 = 64
选择并行策略时通常遵循这些原则：
- 单层太大：优先考虑 TP。
- 层数太多：考虑 PP。
- 训练吞吐不够：增加 DP。
- MoE experts 太多：使用 EP。
- 长序列 activation 显存太高：考虑 SP。
- 通信拓扑很重要：TP 和 EP 更依赖高速互联，跨节点时要谨慎。
11. 推理部署中的 TP 注意点
在大模型推理中，TP 需要重点关注以下问题：
- KV Cache 分布
Attention heads 被切分后，每张 GPU 只保存自己 heads 对应的 K/V Cache。这样 KV Cache 也天然按 head 分片，显存压力随 TP 分摊。
- Prefill 与 Decode 的差异
Prefill 阶段 sequence length 较长，GEMM 和 Attention 计算量大，TP 分摊计算更明显。Decode 阶段通常每次生成一个 token，batch 较小时 AllReduce 延迟更容易影响性能。
- TP size 选择
TP size 越大，单卡显存越小，但通信越多。常见经验是：
- 单卡能放下时，不一定需要 TP。
- 单节点多卡优先选择节点内 TP。
- 不要盲目跨节点做大 TP，除非网络足够快。
- TP size 要和 head 数、hidden size、intermediate size 匹配。
- 通信库和拓扑
TP 强依赖 NCCL / RCCL / mcCCL 等集合通信库，以及 GPU 间拓扑。AllReduce 性能通常决定 TP 扩展效率上限。
12. 简单结论
Tensor Parallel 是当前大模型中最常见、最核心的模型并行方式之一。它把 Transformer 层内部的大矩阵乘法切到多张 GPU 上，典型模式是：
- QKV 和 MLP Linear1 使用 Column Parallel。
- Attention heads 按 GPU 切分，各卡独立计算 attention。
- Output Projection 和 MLP Linear2 使用 Row Parallel。
- Row Parallel 之后通过 AllReduce 合并 partial outputs。
TP 的价值在于让更大的模型能够运行，并把单层计算分摊到多卡；TP 的成本在于每层引入通信。实际部署时，需要在显存、计算吞吐、通信带宽、batch size、sequence length 和模型结构之间做平衡。
更多推荐
6
0- 0
已为社区贡献4条内容
所有评论(0)