系统第 8 讲 / 共 17 讲

第 8 讲：Parallelism（下）

上一讲我们学习了 Data Parallelism 和 ZeRO，它们通过复制模型、分割数据来并行训练。但当模型大到连单层参数都无法装入一个 GPU 时，我们需要将模型本身拆分到多个设备上。本讲将介绍 Tensor Parallelism、Pipeline Parallelism、Sequence Parallelism、Context Parallelism 以及 Expert Parallelism，最后讨论如何将这些策略组合成 3D Parallelism。

Tensor Parallelism (TP)

Tensor Parallelism 的核心思想是将单个层的权重矩阵切分到多个 GPU 上，每个 GPU 只计算矩阵乘法的一部分，最后通过通信合并结果。这是一种模型内并行（intra-layer parallelism）。

Column Parallel Linear

考虑一个线性层 Y = XA，其中 X ∈ R^{b×k}，A ∈ R^{k×n}。Column Parallel 将权重 A 按列切分：

A = [A₁ | A₂ | ... | A_W]，其中 A_i ∈ R^{k×(n/W)}

每个 GPU i 持有 A_i，并计算：

Y_i = X × A_i ∈ R^{b×(n/W)}

结果 Y_i 是输出矩阵 Y 的一部分列。所有 Y_i 拼接起来（concatenate）就得到完整的 Y。

通信需求：输入 X 需要被复制到（或已经在）所有 GPU 上。输出可以保持分片状态，传递给下一层使用。

Row Parallel Linear

Row Parallel 将权重 A 按行切分：

A = [A₁; A₂; ...; A_W]，其中 A_i ∈ R^{(k/W)×n}

对应地，输入 X 也需要按列切分：X_i ∈ R^{b×(k/W)}。每个 GPU 计算：

Y_i = X_i × A_i ∈ R^{b×n}

最终结果需要 AllReduce（求和）：

Y = Σ_{i=1}^{W} Y_i

通信需求：输入已经是分片的（每个 GPU 有不同的列），输出需要 AllReduce 来聚合。

Megatron-LM 风格的 Tensor Parallelism

Megatron-LM（NVIDIA）巧妙地将 Column Parallel 和 Row Parallel 组合起来，最小化 Transformer 中的通信次数。

MLP 层的 TP 策略：

Transformer 的 MLP 层通常是两个线性变换夹一个激活函数：

MLP(x) = Dropout(GeLU(xA) × B)

Megatron-LM 的做法：

第一个线性层 A 用 Column Parallel：每个 GPU 得到 GeLU 输出的一部分列
第二个线性层 B 用 Row Parallel：输入刚好是分片的，AllReduce 输出

这样整个 MLP 只需要 1 次 AllReduce（在 forward pass 中）。

Attention 层的 TP 策略：

Multi-Head Attention 天然适合 Tensor Parallelism——不同的 attention head 可以分配到不同的 GPU。

Q、K、V 的投影矩阵用 Column Parallel：每个 GPU 负责一部分 head
输出投影矩阵用 Row Parallel：AllReduce 合并结果

同样，整个 Attention 层在 forward 中只需要 1 次 AllReduce。

Megatron-LM 的通信效率

每个 Transformer layer 在 forward pass 中只需要 2 次 AllReduce（MLP 1 次 + Attention 1 次），backward pass 也是 2 次。考虑到 AllReduce 可以与计算 overlap，TP 的通信开销是可控的。但 AllReduce 的 latency 要求 TP 通常只在节点内高带宽互连（如 NVLink）的 GPU 之间使用。

# Pseudo-code for Column Parallel Linear (forward)
class ColumnParallelLinear(nn.Module):
    def __init__(self, in_features, out_features, tp_size):
        super().__init__()
        self.tp_size = tp_size
        self.tp_rank = dist.get_rank()
        # Each GPU has out_features // tp_size columns
        self.weight = nn.Parameter(
            torch.empty(in_features, out_features // tp_size)
        )

    def forward(self, x):
        # x: [batch, seq_len, in_features] - replicated on all GPUs
        # output: [batch, seq_len, out_features // tp_size] - sharded
        return F.linear(x, self.weight.T)

# Pseudo-code for Row Parallel Linear (forward)
class RowParallelLinear(nn.Module):
    def __init__(self, in_features, out_features, tp_size):
        super().__init__()
        self.tp_size = tp_size
        # Each GPU has in_features // tp_size rows
        self.weight = nn.Parameter(
            torch.empty(in_features // tp_size, out_features)
        )

    def forward(self, x):
        # x: [batch, seq_len, in_features // tp_size] - sharded input
        local_out = F.linear(x, self.weight.T)
        # AllReduce to sum partial results
        dist.all_reduce(local_out, op=dist.ReduceOp.SUM)
        return local_out

Pipeline Parallelism (PP)

Pipeline Parallelism 将模型按层分割到不同 GPU 上。例如一个 32 层的 Transformer，可以将前 8 层放在 GPU 0，第 9-16 层放在 GPU 1，以此类推。每个 GPU 只执行模型的一个"阶段（stage）"。

Naive Pipeline 的问题

最朴素的实现是：GPU 0 完成所有 micro-batch 的前 8 层计算后，才将结果传给 GPU 1。这导致了严重的 Pipeline Bubble（流水线气泡）——大部分时间里，大部分 GPU 都在空闲等待。

Bubble Ratio = (P - 1) / (P - 1 + M)

其中 P 是 pipeline stages 数，M 是 micro-batches 数。当 P = 4, M = 1 时，bubble 高达 75%——4 个 GPU 的效率还不如 1 个。

GPipe Schedule

GPipe（Google, 2019）通过 Micro-batching 来减少 bubble。将一个 mini-batch 分成 M 个 micro-batch，让各 stage 像流水线一样处理：

所有 micro-batch 先完成 forward pass（从 stage 0 到 stage P-1）
所有 micro-batch 再完成 backward pass（从 stage P-1 到 stage 0）

当 M >> P 时，bubble ratio 趋近于 0。但 GPipe 要求同时存储所有 micro-batch 的 activation，内存占用较高。

1F1B Schedule（Interleaved）

1F1B（One Forward One Backward）是 PipeDream 提出的更高效的调度方案：

Warmup 阶段：各 stage 依次启动 forward，填满 pipeline
Steady State：每个 stage 交替执行一个 forward 和一个 backward
Cooldown 阶段：清空 pipeline 中剩余的 backward

1F1B 的优势是内存效率高——每个 stage 同时只需要保存少量 micro-batch 的 activation（最多 P 个），而 GPipe 需要保存所有 M 个。

Pipeline Bubble 的实际影响

在实际训练中，通常设置 M ≥ 4P 来将 bubble overhead 控制在可接受的范围内（< 20%）。例如 4 个 pipeline stages，至少需要 16 个 micro-batches。这意味着 PP 通常与较大的 batch size 配合使用。

Pipeline Parallelism 的通信特点

PP 的通信模式与 TP 截然不同：

通信量小：只需传递 activation tensors（不是完整的参数或梯度），大小为 b × s × h（batch × seq_len × hidden_size）
点对点通信：只在相邻 stage 之间传输，不需要 collective communication
对带宽要求低：适合跨节点使用（通过 InfiniBand）

Sequence Parallelism (SP)

在 Tensor Parallelism 中，LayerNorm 和 Dropout 这些操作是在每个 GPU 上重复执行的（因为它们需要完整的 hidden dimension）。Sequence Parallelism 将这些操作沿 sequence 维度切分，进一步减少冗余计算和内存占用。

具体来说：

在 TP 区域（Attention、MLP 的矩阵乘法部分），数据沿 hidden dimension 切分
在非 TP 区域（LayerNorm、Dropout），数据沿 sequence dimension 切分
在两种区域之间，通过 AllGather / ReduceScatter 进行转换

Sequence Parallelism 的核心收益是减少 activation memory。对于一个长度为 s 的序列，每个 GPU 只需要存储 s/W 部分的 activation。

Activation Memory Reduction = W× (for LayerNorm and Dropout activations)

Context Parallelism (Ring Attention)

当序列长度极长时（如 128K、1M tokens），即使是单个序列的 attention 计算也会超出 GPU 显存。Context Parallelism（也叫 Ring Attention）将一个长序列切分到多个 GPU 上，每个 GPU 只处理序列的一部分。

Ring Attention 的工作流程

将长序列切成 W 块，分配给 W 个 GPU
每个 GPU 有自己的 Q（query）块，需要与所有 K、V 块做 attention
GPU 排成环形：每个 GPU 持有一块 K、V，计算完 partial attention 后，将 K、V 传给环中的下一个 GPU
经过 W-1 轮传递后，每个 GPU 的 Q 都与所有 K、V 完成了 attention 计算

Ring Attention 的关键在于通信与计算的 overlap：当 GPU 在计算当前 K、V 块的 attention 时，同时在传输下一轮需要的 K、V 块。

Context Parallelism vs Sequence Parallelism

虽然两者都涉及序列的切分，但它们的目标不同：Sequence Parallelism 是 Tensor Parallelism 的补充，减少 LayerNorm/Dropout 的 activation 冗余；Context Parallelism 是为了处理超长序列，将 attention 计算本身分布到多个 GPU 上。

Expert Parallelism (EP)

在 Mixture of Experts (MoE) 架构中，MLP 层被替换为多个 expert 子网络，每个 token 只激活少量 expert。Expert Parallelism 将不同的 expert 放在不同的 GPU 上。

EP 的通信模式

EP 使用 All-to-All 通信原语：

Dispatch：Router 决定每个 token 由哪些 expert 处理后，通过 All-to-All 将 token 发送到对应 expert 所在的 GPU
Compute：每个 GPU 处理分配到自己 expert 的 tokens
Combine：通过 All-to-All 将结果发送回 token 原来所在的 GPU

All-to-All 的通信模式是不规则的——每对 GPU 之间交换的数据量取决于 router 的决策。这使得 load balancing 成为 MoE 训练的关键挑战。

# EP communication pattern (simplified)
def expert_parallel_forward(tokens, router, experts, ep_group):
    # Step 1: Route tokens
    routing_weights, expert_indices = router(tokens)

    # Step 2: All-to-All dispatch
    # Send tokens to the GPU that owns their assigned expert
    dispatched_tokens = all_to_all(tokens, expert_indices, group=ep_group)

    # Step 3: Local expert computation
    expert_output = experts[local_expert_id](dispatched_tokens)

    # Step 4: All-to-All combine
    # Send results back to original GPU
    combined_output = all_to_all(expert_output, reverse_indices, group=ep_group)

    return combined_output * routing_weights

3D Parallelism 组合策略

在训练超大模型时，通常需要同时使用多种并行策略。最经典的组合是 3D Parallelism：Data Parallelism × Tensor Parallelism × Pipeline Parallelism。

设计原则

不同的并行策略有不同的通信特征，需要合理匹配到硬件拓扑：

并行策略	通信模式	通信量	Latency 敏感度	推荐互连
Tensor Parallelism	AllReduce (频繁)	高	高	NVLink (节点内)
Pipeline Parallelism	P2P (顺序)	低	中	InfiniBand (跨节点)
Data Parallelism	AllReduce (每步一次)	高	低	InfiniBand (跨节点)
Expert Parallelism	All-to-All	中	中	NVLink 或 IB

典型配置方案

以一个包含 512 个 GPU（64 节点，每节点 8 GPU）的集群为例：

配置示例：训练 175B 模型 (类 GPT-3)

TP = 8（节点内 8 个 GPU 用 NVLink 做 Tensor Parallelism）
PP = 8（8 个节点串成 pipeline，用 InfiniBand）
DP = 8（8 组 pipeline replica，用 InfiniBand AllReduce）
总 GPU 数 = 8 × 8 × 8 = 512

在 3D Parallelism 中，GPU 被组织成一个三维网格：

TP 维度（最内层）：同一节点内的 GPU，通过 NVLink 高速互连
PP 维度（中间层）：跨节点，但通信量小，每次只传递 activation
DP 维度（最外层）：跨节点，梯度同步可以与计算 overlap

实际训练配置案例

模型	参数量	TP	PP	DP	总 GPU	硬件
GPT-3	175B	8	8	8	512	V100 32GB
Megatron-Turing NLG	530B	8	35	-	560	A100 80GB
LLaMA-65B	65B	8	-	256	2048	A100 80GB
LLaMA-3 405B	405B	8	16	-	16384	H100 80GB

为什么 LLaMA-65B 不用 Pipeline Parallelism？

LLaMA-65B "只有" 65B 参数，使用 TP=8 后每个 GPU 约 8B 参数，可以装入 A100 80GB。此时不需要 PP，而是用 ZeRO + DP 就足够了。PP 引入的 pipeline bubble 会降低效率，只有当模型大到 TP+ZeRO 都不够时才值得使用。

更先进的并行策略

Interleaved Pipeline Parallelism

Megatron-LM 还提出了 Interleaved Pipeline Schedule：不再将连续的层分配到一个 stage，而是让每个 stage 负责多个不连续的层组（virtual stages）。例如 GPU 0 负责层 0-3 和层 16-19，GPU 1 负责层 4-7 和层 20-23。

这样可以将 pipeline bubble 减半，代价是增加了通信次数（每个 virtual stage 之间都需要 P2P 通信）。

Zero Bubble Pipeline

2024 年的研究提出了接近零 bubble 的 pipeline schedule。核心思想是将 backward pass 拆分为两部分：

B（backward for input gradient）：计算输入梯度，传给上一个 stage
W（backward for weight gradient）：计算权重梯度，用于参数更新

通过灵活调度 F、B、W 操作，可以在 steady state 几乎消除所有 bubble。

小结

Tensor Parallelism：在层内切分权重矩阵，通过 Column/Row Parallel 实现，通信频繁但数据量可控，适合节点内 NVLink
Pipeline Parallelism：按层切分模型，使用 micro-batching 减少 bubble（GPipe 或 1F1B schedule），适合跨节点
Sequence Parallelism：TP 的补充，将 LayerNorm/Dropout 沿序列维度分片
Context Parallelism (Ring Attention)：处理超长序列，将 attention 分布到多 GPU
Expert Parallelism：MoE 模型中将不同 expert 放在不同 GPU，使用 All-to-All 通信
3D Parallelism：TP（节点内）× PP（跨节点）× DP（跨节点组）的组合，是训练超大模型的标准方案