LLM Parallelism 基础

Tensor Parallelism

MPI Collective Operations

MPI 定义了一组集合通信操作（collective operations）——让多个进程协调通信。AllReduce 就是其中之一。

AllReduce

字面拆解：

Reduce = 把多个值"归约"成一个（求和、求最大、求最小等）
All = 所有进程都得到这个归约结果

操作语义：

初始状态（4 个进程）：
  进程 0: 值 = a
  进程 1: 值 = b
  进程 2: 值 = c
  进程 3: 值 = d

AllReduce(SUM) 之后：
  进程 0: 值 = a + b + c + d
  进程 1: 值 = a + b + c + d
  进程 2: 值 = a + b + c + d
  进程 3: 值 = a + b + c + d

所有进程最终都拥有同一个全局求和结果。

其他 Collective

操作	含义	例子
Broadcast	一个进程发，所有进程收	rank 0 把数据发给所有人
Gather	所有进程的数据汇总到一个进程	4 个进程的结果都送到 rank 0
Scatter	一个进程把数据切片发给所有进程	rank 0 把数组切成 4 份发给 4 个进程
AllGather	所有进程的数据汇总到所有进程	每个进程都得到所有进程的数据
Reduce	所有进程归约后送到一个进程	求和后只有 rank 0 拿到结果
AllReduce	所有进程归约后所有人都拿到	求和后所有进程都拿到结果 ⭐
ReduceScatter	归约 + 切片分发	AllReduce 的"省一半带宽"版本

注意——AllReduce = Reduce + Broadcast 的组合效果——但实际实现可以更高效（不是真的串行做两步）。

Ring AllReduce

AllReduce 是分布式数据并行训练（Data Parallelism）中绝对的核心——每张显卡算出了自己那部分数据的梯度，必须通过 AllReduce 将所有卡的梯度加起来，然后再分发给每张卡，以此更新模型权重。

如果采用 Naive 的方式（大家同时互相发数据），网络会瞬间拥塞。NCCL 为了解决这个问题，采用了针对特定拓扑高度优化的算法。最经典且最核心的算法是 Ring-AllReduce（环形 AllReduce）。

为了最大化带宽利用率，Ring-AllReduce 将所有参与的 GPU 连成一个逻辑上的“环”，并将你要传输的数据（比如一个巨大的浮点数组）平均切分成 $N$ 块（$N$ 为 GPU 的数量）。

整个过程被巧妙地拆分为两个阶段，总共需要 $2 \times (N - 1)$ 步：

阶段 1：Reduce-Scatter (规约-散发)

目标：让每张显卡最终只持有完整规约结果的 $1 / N$ 。

在每一步中，每张卡只把自己的一小块数据发给环上的下一张卡。
下一张卡收到后，与自己对应的本地数据块进行累加（规约）。
经过 $N - 1$ 步流转，整个数组的每一块都在环上转了一圈，完成了所有卡的累加。此时，卡 0 持有第 0 块的最终结果，卡 1 持有第 1 块的最终结果……

阶段 2：All-Gather (全收集)

目标：将大家手里那 $1 / N$ 的完整结果，分享给所有人。

还是沿着那个环，每张卡把自己拥有的“最终结果块”发给下一张卡。
下一张卡收到后，直接覆盖本地对应的未完全规约的数据块，并继续向后传递。
再经过 $N - 1$ 步，所有显卡都拼凑出了完整的、规约后的最终数组。

在这个过程中，无论有多少张显卡参与，每张卡在每一步的接收量和发送量都是恒定不变的。系统没有网络热点，没有通信瓶颈，完美地吃满了环上任意两点之间的链路带宽。

矩阵乘法 XW 分割

假设要算 $Y = X \cdot W$ ——

X: (batch, d_in)
W: (d_in, d_out)
Y: (batch, d_out)

两张卡（TP=2）怎么切？

方式 A：按列切 W（Column Parallel）

把 $W$ 沿着列方向切成两半——

W = [W_1 | W_2]
     ↑     ↑
   (d_in,    (d_in,
   d_out/2) d_out/2)

GPU 0 持有 W_1
GPU 1 持有 W_2

每张卡算各自的部分：

GPU 0: Y_1 = X · W_1   shape (batch, d_out/2)
GPU 1: Y_2 = X · W_2   shape (batch, d_out/2)

最后拼接两张卡的结果（concat）：

Y = [Y_1 | Y_2]    shape (batch, d_out)

方式 B：按行切 W（Row Parallel）

把 $W$ 沿着行方向切成两半——

W = [ W_1 ]      W_1: (d_in/2, d_out)
    [ W_2 ]      W_2: (d_in/2, d_out)

GPU 0 持有 W_1
GPU 1 持有 W_2

但是这样切——X 也得切成两半：

X = [X_1 | X_2]    X_1: (batch, d_in/2)
                   X_2: (batch, d_in/2)

每张卡算——

GPU 0: Y_1 = X_1 · W_1   shape (batch, d_out)
GPU 1: Y_2 = X_2 · W_2   shape (batch, d_out)

最后两个结果相加（all-reduce sum）——

Y = Y_1 + Y_2    shape (batch, d_out)

两种切法的对比

切法	W 怎么切	输入 X 怎么处理	输出怎么合并
Column Parallel	按列切（垂直切）	X 完整传给每张卡	concat
Row Parallel	按行切（水平切）	X 也按列切	sum（all-reduce）

记住这两个模式——下面看 Megatron 怎么用。

Megatron 风格的 TP

Megatron-LM（Nvidia 2019）是最经典的 TP 切分方案——LLaMA、GPT-3 等都用这个思路。

Attention 怎么切

attention 的关键观察——不同的 head 之间是完全独立的——所以可以按 head 切。

LLaMA 2 7B 有 32 个 head——TP=2 时——

GPU 0 持有 head 0-15
GPU 1 持有 head 16-31

具体来说——

Step 1：W_Q, W_K, W_V 按列切（按 head 切，本质上还是 Column Parallel）——

W_Q: (4096, 4096) → 每张卡持有 (4096, 2048)
W_K: (4096, 4096) → 每张卡持有 (4096, 2048)
W_V: (4096, 4096) → 每张卡持有 (4096, 2048)

每张卡输入完整的 x，算自己那 16 个 head 的 Q/K/V——

GPU 0: Q_0 = x · W_Q_0   shape (batch, seq, 2048)  → reshape (batch, 16 heads, seq, 128)
       K_0, V_0 同理

GPU 1: Q_1 = x · W_Q_1   shape (batch, seq, 2048)  → reshape (batch, 16 heads, seq, 128)
       K_1, V_1 同理

Step 2：每张卡独立算自己那 16 个 head 的 attention——

GPU 0: head 0-15 的 attention 输出   shape (batch, 16 heads, seq, 128)
GPU 1: head 16-31 的 attention 输出  shape (batch, 16 heads, seq, 128)

这一步不需要通信——不同 head 之间不互相访问。

Step 3：W_O 按行切（Row Parallel）——

W_O: (4096, 4096) → 切成 [ W_O_0 ]   shape (2048, 4096)
                          [ W_O_1 ]   shape (2048, 4096)

每张卡算自己那一半的输出投影——

GPU 0: out_0 = attention_output_0 · W_O_0   shape (batch, seq, 4096)
GPU 1: out_1 = attention_output_1 · W_O_1   shape (batch, seq, 4096)

Step 4：All-reduce sum——

attention_output = all_reduce_sum(out_0, out_1)   shape (batch, seq, 4096)

Attention 切分的关键观察

为什么按 head 切是天然的？

因为 multi-head attention 的设计本身就是 h 个独立的 head 并行计算 + 最后用 W_O 融合——

每个 head 是独立的 attention 计算（不互相访问）
W_O 把 h 个 head 的结果融合

TP 完美对应这个结构——

每张卡处理一部分 head（独立计算）
W_O 行切 + all-reduce 完成"融合"

这是为什么 TP 切 attention 几乎没有"概念上的损失"——它是在沿模型架构本身的天然边界切。

FFN 怎么切

LLaMA 的 FFN（SwiGLU）：

x: (batch, seq, 4096)
   │
   ├── × W_gate (4096, 11008) → SiLU
   ├── × W_up (4096, 11008)
   │       ↓
   │     逐元素相乘
   │       ↓
   └── × W_down (11008, 4096) → output

Megatron 的切法：

Step 1：W_gate 和 W_up 按列切（Column Parallel）——

W_gate: (4096, 11008) → 切成 [W_gate_1 | W_gate_2]
                        每张卡持有 (4096, 5504)

W_up:   (4096, 11008) → 切成 [W_up_1 | W_up_2]
                        每张卡持有 (4096, 5504)

每张卡输入完整的 x，算自己那一半——

GPU 0: gate_1 = x · W_gate_1   shape (batch, seq, 5504)
       up_1   = x · W_up_1     shape (batch, seq, 5504)
       activated_1 = SiLU(gate_1) * up_1   shape (batch, seq, 5504)

GPU 1: gate_2 = x · W_gate_2   shape (batch, seq, 5504)
       up_2   = x · W_up_2     shape (batch, seq, 5504)
       activated_2 = SiLU(gate_2) * up_2   shape (batch, seq, 5504)

这一步不需要通信——每张卡独立算自己那一半。

Step 2：W_down 按行切（Row Parallel）——

W_down: (11008, 4096) → 切成 [ W_down_1 ]   shape (5504, 4096)
                              [ W_down_2 ]   shape (5504, 4096)

每张卡算自己那一半——

GPU 0: out_1 = activated_1 · W_down_1   shape (batch, seq, 4096)
GPU 1: out_2 = activated_2 · W_down_2   shape (batch, seq, 4096)

Step 3：All-reduce sum——

output = all_reduce_sum(out_1, out_2)   shape (batch, seq, 4096)
        = out_1 + out_2 (跨 GPU 求和)

FFN 这么切的妙处

为什么是"先 Column Parallel 再 Row Parallel"？

关键——这个组合让 FFN 内部不需要通信——只在 FFN 出口做一次 all-reduce。

具体看——

x → [Column Parallel: W_gate, W_up]     ← 不通信
       ↓
     SiLU + 逐元素相乘                     ← 不通信
       ↓
     [Row Parallel: W_down]              ← 算完后才 all-reduce
       ↓
     all-reduce sum                       ← 通信

整个 FFN 只有最后一次通信——通信开销最小。

如果反过来（先 Row Parallel 再 Column Parallel）——FFN 中间就要通信——开销大得多。

完整的 LLaMA Block 在 TP=2 下的数据流

让我把整个 decoder block 在 TP=2 下的数据流画出来——

x: (batch, seq, 4096)   ← 每张卡都有完整的 x
       │
       ↓
   RMSNorm                ← 每张卡独立算（参数小，每张卡复制一份）
       │
       ↓
   ┌───────────────────────────────────────────────────┐
   │ Attention                                         │
   │                                                   │
   │ Q, K, V (按 head 切)         ← 不通信              │
   │   GPU 0: Q,K,V for head 0-15                      │
   │   GPU 1: Q,K,V for head 16-31                     │
   │                                                   │
   │ 各自算 attention             ← 不通信               │
   │   GPU 0: attention_0 (16 heads)                   │
   │   GPU 1: attention_1 (16 heads)                   │
   │                                                   │
   │ W_O (Row Parallel)           ← 算完后 all-reduce   │
   │   ALL-REDUCE SUM →           ← 通信发生在这里！      │
   │ output: (batch, seq, 4096)                        │
   └───────────────────────────────────────────────────┘
       │
       ↓
   + 残差
       │
       ↓
   RMSNorm                 ← 每张卡独立算
       │
       ↓
   ┌───────────────────────────────────────────────────┐
   │ FFN (SwiGLU)                                      │
   │                                                   │
   │ W_gate, W_up (Column Parallel) ← 不通信            │
   │   GPU 0: gate_0, up_0 (5504 维)                   │
   │   GPU 1: gate_1, up_1 (5504 维)                   │
   │                                                   │
   │ SiLU + 逐元素相乘             ← 不通信               │
   │                                                   │
   │ W_down (Row Parallel)         ← 算完后 all-reduce  │
   │   ALL-REDUCE SUM →            ← 通信发生在这里！     │
   │ output: (batch, seq, 4096)                        │
   └───────────────────────────────────────────────────┘
       │
       ↓
   + 残差
       │
       ↓
   下一层 block

注意——每个 block 有 2 次 all-reduce——一次在 attention 后，一次在 FFN 后。

LLaMA 2 7B 有 32 层——所以一次 forward pass 总共有 64 次 all-reduce。

通信开销

每次 all-reduce 通信的数据量——

all-reduce 的数据量 = (batch × seq × 4096) × dtype_size

对于 batch=1, seq=2048, FP16——

通信量 = 1 × 2048 × 4096 × 2 字节 = 16 MB / 次

64 次 all-reduce——总通信 ≈ 1 GB（per forward pass）。

NVLink 带宽 600 GB/s——理论通信时间约 1.7 ms（per forward pass）——

但实际上——

每次 all-reduce 有固定的 launch overhead
Decode 阶段每次只算 1 个 token——通信开销在总时间里占比比 prefill 高得多
Inter-node（跨服务器）的通信带宽（200 Gbps Infiniband）远低于 intra-node（NVLink）

这就是为什么 TP 的"sweet spot"通常是单节点内 8 卡——跨节点会被网络拖死。

一个直观的对比

让我用数字感受一下 TP 的效果——

没有 TP（单卡运行 LLaMA 2 7B）：

模型权重 14 GB——必须在一张卡的显存里
每步推理读 14 GB 权重 + KV cache
A100 80GB 能跑

TP=2：

每张卡持有 7 GB 权重
每步推理每张卡读 7 GB 权重——理论上 decode 速度 ×2
但有 64 次 all-reduce 通信开销——实际加速 < 2x
A100 40GB 也能跑了

TP=8：

每张卡持有 1.75 GB 权重
每步推理每张卡读 1.75 GB 权重——理论上 decode 速度 ×8
通信开销更大——实际加速可能只有 4-5x
但能跑更大的模型（LLaMA 70B 单卡跑不动，TP=8 能跑）

TP 的本质

TP 在做什么？

把"一个大矩阵乘法"拆成"多张卡协作的多个小矩阵乘法"——每张卡持有一部分参数、做一部分计算——通过 all-reduce 合并结果。

它解决的核心问题——

单卡显存不够：把模型参数分散到多卡（70B 模型单卡装不下）
单卡算力不够：让多卡并行算（每张卡只做 1/N 的工作）

代价——

通信开销（all-reduce）
实现复杂度（需要 NCCL 等通信库）
卡数越多，通信开销占比越大——并行效率递减🎯