DDP 解决的是样本维度扩展:每张 GPU 都有完整模型,只是处理不同数据。 ZeRO/FSDP 解决的是模型状态切分:参数、梯度、优化器状态不再完整复制。 Tensor Parallel 解决的是另一类问题: 单层矩阵本身太大,或单个 GEMM 太重,需要多张 GPU 一起算同一层。 对 Transformer 来说,最值得切的是 MLP 的两层大 Linear、Attention 的 QKV projection 和 output projection。

TP 最适合放在哪里

Megatron-LM 的 TP 可以先抽象成一个矩阵乘法: Y = X A + b。 如果把 A 沿输出维度切,每张 GPU 算一段 Y_i = X A_i,这就是列并行。 如果把 A 沿输入维度切,输入 X 也被切成 X_i,每张 GPU 算 partial sum, 最后 reduce 成完整输出,这就是行并行。

Megatron 两种 Linear 的语义

# 伪代码:Megatron-style MLP 的核心形状
# X: [S, B, H], tensor_model_parallel_size = tp

# 1) 列并行:每张 GPU 算 4H/tp 个 hidden
Y_i = X @ W1_i + b1_i                 # W1_i: [H, 4H/tp]
Z_i = gelu(Y_i)                       # activation 在本地做,不通信

# 2) 行并行:输入已经是 Z_i,每张 GPU 算输出 partial sum
P_i = Z_i @ W2_i                      # W2_i: [4H/tp, H]
O = all_reduce_sum(P_i) + b2          # O: [S, B, H]

以标准 FFN/MLP 为例: MLP(X) = GELU(X W_up) W_down。 Megatron 的设计很漂亮:第一层 W_up 做列并行, 输出的 4H 中间维度被切成多份; GELU 是逐元素函数,每张 GPU 在本地做; 第二层 W_down 做行并行,用本地中间激活算 partial output, 最后一次 AllReduce 得到完整 hidden。

形状表:以 TP=4 为例

Multi-Head Attention 比 MLP 更自然适合 TP: Q/K/V projection 的输出可以按 heads 切分。 如果总共有 num_heads 个头,TP size 是 p,每个 rank 持有 num_heads / p 个头。 每个 rank 在本地完成这些头的 attention softmax 和 value 加权, 然后 output projection 用 RowParallelLinear 把各 rank 的 head 输出合成完整 hidden。

Attention 的通信边界

# 伪代码:Megatron-style attention TP
# X: [S, B, H], num_heads = A, tp = p

QKV_i = X @ W_qkv_i                   # ColumnParallel, [S,B,3H/p]
Q_i, K_i, V_i = split_heads(QKV_i)    # each rank owns A/p heads

Context_i = attention(Q_i, K_i, V_i)  # local heads only, no TP communication
P_i = Context_i @ W_o_i               # RowParallel partial output
O = all_reduce_sum(P_i) + b_o         # full [S,B,H]

TP 看起来很美:参数和 GEMM 工作量约按 1 / tp 分到每张 GPU。 但通信也进入了每个 Transformer layer 内部。 所以 TP size 不是越大越好:太小放不下层,太大每张 GPU 的 GEMM 变碎、通信比例上升、kernel 利用率下降。 这也是为什么实际大模型常把 TP 放在单机高速互联内,再用 PP/DP 扩到多机。

选择 TP size 的检查表

今天不要求完整跑 Megatron 训练,那是 Day22 的任务。 今天的重点是能从配置推导出每个 rank 持有哪些权重分片、哪些激活分片、哪些通信边界。 你可以用纸画,也可以把下面模板填成 markdown 表。

Megatron 关键配置

# 典型 Megatron-LM 配置片段
torchrun --nproc_per_node=8 pretrain_gpt.py \
  --tensor-model-parallel-size 4 \
  --pipeline-model-parallel-size 2 \
  --num-layers 32 \
  --hidden-size 4096 \
  --ffn-hidden-size 16384 \
  --num-attention-heads 32 \
  --sequence-parallel \
  ...

# 解释:
# 8 GPUs = TP 4 * PP 2 * DP 1
# 每个 pipeline stage 内有一个 4-GPU tensor parallel group
# 每个 TP rank 持有 32/4 = 8 个 attention heads

推导模板

小型教学模拟:不用 Megatron 也能理解

# 这是单进程 shape 模拟,帮助你检查列切/行切的数学等价性
import torch

torch.manual_seed(0)
B, S, H, TP = 2, 3, 8, 2
X = torch.randn(B, S, H)
W1 = torch.randn(H, 4 * H)
W2 = torch.randn(4 * H, H)

# eager full MLP
full = torch.relu(X @ W1) @ W2

# ColumnParallel W1: split output dim
W1_parts = W1.chunk(TP, dim=1)
Z_parts = [torch.relu(X @ w) for w in W1_parts]

# RowParallel W2: split input dim, then reduce partial sums
W2_parts = W2.chunk(TP, dim=0)
partials = [z @ w for z, w in zip(Z_parts, W2_parts)]
tp_out = sum(partials)                 # distributed version = all_reduce_sum

torch.testing.assert_close(tp_out, full)
print("TP math matches full MLP")

交付物检查清单

1. Tensor Parallel 和 Data Parallel 分别切的是哪个维度?
2. ColumnParallelLinear 为什么切输出维度?它什么时候需要 all-gather?
3. RowParallelLinear 为什么需要 reduce partial sums?
4. Megatron MLP 为什么可以把中间激活保持为分片状态?
5. Attention 的 QKV projection 为什么可以按 heads 切?
6. 给定 num_heads=40,哪些 TP size 是合法的?为什么?
7. TP size 过大时,性能可能从哪里掉下来?