当前位置: 首页 > news >正文

大模型中的各种并行: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=2TP=4TP=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 通常使用两类切分方式:

  1. Column Parallel:按输出维度切分权重列。
  2. 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 ProjectionColumn Parallel / Head Split部分 QKV heads通常不需要
Attention 计算Head Parallel部分 heads 的 attention通常不需要
Output ProjectionRow Parallel部分输入行权重需要 AllReduce
MLP Linear1Column Parallel部分 intermediate 输出列通常不需要立即 Gather
Activation / GateLocal本卡 intermediate 分片不需要
MLP Linear2Row Parallel部分 intermediate 输入行需要 AllReduce
LayerNorm / ResidualReplicated Hidden完整 hidden通常依赖前序 AllReduce 后本地执行

一个标准 Transformer Decoder Layer 中,TP 通常会产生两次主要 AllReduce:

  1. Attention Output Projection 之后。
  2. 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=8TP=4。跨节点扩展时,通常会组合 PP、DP、EP 等其他并行方式。

10. 大模型中的其他常见并行方式概括

前面重点介绍了 TP。真实大模型训练和推理系统通常不会只使用 TP,而是把 TP、DP、PP、EP、SP 等并行方式组合起来。不同并行方式切分的对象不同,解决的问题也不同。

并行方式全称主要切分对象主要解决的问题典型通信
TPTensor Parallel单层内部的矩阵、hidden 维度、attention heads单层太大,单卡放不下或算不动AllReduce / AllGather / ReduceScatter
DPData Parallelbatch 数据多副本训练,提高训练吞吐梯度 AllReduce
PPPipeline ParallelTransformer 层模型层数太多,无法放在单卡或单节点stage 间 activation / gradient 传递
EPExpert ParallelMoE expertsMoE expert 数量多,单卡放不下全部 expertstoken dispatch / combine,AllToAll
SPSequence Parallelsequence 维度上的激活降低长序列训练中的激活显存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 需要重点关注以下问题:

  1. KV Cache 分布

Attention heads 被切分后,每张 GPU 只保存自己 heads 对应的 K/V Cache。这样 KV Cache 也天然按 head 分片,显存压力随 TP 分摊。

  1. Prefill 与 Decode 的差异

Prefill 阶段 sequence length 较长,GEMM 和 Attention 计算量大,TP 分摊计算更明显。Decode 阶段通常每次生成一个 token,batch 较小时 AllReduce 延迟更容易影响性能。

  1. TP size 选择

TP size 越大,单卡显存越小,但通信越多。常见经验是:

  • 单卡能放下时,不一定需要 TP。
  • 单节点多卡优先选择节点内 TP。
  • 不要盲目跨节点做大 TP,除非网络足够快。
  • TP size 要和 head 数、hidden size、intermediate size 匹配。
  1. 通信库和拓扑

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 和模型结构之间做平衡。

http://www.cnnetsun.cn/news/3155806.html

相关文章:

  • 鸿蒙 CodeGenie:技能(Skills)配置
  • openEuler-pkginfo错误排查指南:常见问题与解决方案
  • WhatsApp 自动回复规则引擎的设计与实现
  • openEuler-pkginfo性能优化:如何高效处理大规模仓库数据的10个技巧
  • openEuler-pkginfo扩展开发:5个步骤轻松添加自定义功能模块
  • openeuler/cloudphone_kernel 常见问题解答:新手必看的10个实用技巧
  • openeuler/riscv-kernel项目架构深度解析:如何实现多SoC平台统一支持
  • 08_检查点
  • AI驱动的钱包交易风险解释:让链上操作在签名前可理解
  • IIM-42652 IMU传感器与STM32的6DoF运动追踪实现
  • openeuler/riscv-kernel测试与验证:确保内核稳定性的完整方法
  • 如何快速配置Autovisor:完整智慧树刷课脚本使用教程
  • AI Agent:自主智能体的工作原理与应用全景
  • ICM-42688-P与PIC18F85J50在运动控制与振动监测中的应用
  • 工业传感器控制系统核心组件与接口设计指南
  • AI 创作时间线:灵感、草稿和成稿不要混在一起
  • migration-assistant源码解析:核心模块架构与代码实现
  • ncmdump:3步解锁网易云音乐加密格式,让音乐真正属于你!
  • 工业级传感器控制系统设计与芯片选型指南
  • 基于FOC算法的15A无刷电机控制方案解析
  • BLE通信安全实践:基于AES128的加密实现与协议栈解析
  • OpenCV环境搭建与人脸识别实战:从零到一的Python计算机视觉入门
  • 直流有刷电机高效驱动方案与PIC微控制器应用
  • TM4C129ENCPDT微控制器与SLO2016协议栈应用解析
  • DS28EC20 EEPROM与PIC18F26K22微控制器的嵌入式存储方案
  • 给 Claude Code 省 97% Token 是真的吗?我把 caveman 装上跑了一周
  • TensorFlow Lite Micro 算子裁剪:少注册一个算子,省半块 Flash
  • 基于Si4731与MK64FX512VDC12的广播接收系统设计
  • 13DOF传感器与PIC18F27J53在AGV导航中的应用
  • STM32与EEPROM嵌入式存储方案设计与实现