第14章张量并行与流水线并行

"Divide and conquer, but make sure the pieces can talk to each other."

本章要点

理解张量并行（TP）的原理：如何在不改变数学结果的前提下切分矩阵乘法
掌握 Megatron-LM 张量并行方案在 vLLM 中的应用
理解流水线并行（PP）的基本原理与调度策略
深入 NCCL 通信原语在 vLLM 中的使用
认识 TP vs PP 的选择策略：什么时候用哪种

14.1 大模型的困境

Llama-2-70B 在 FP16 下需要 140 GB 显存。一张 A100-80GB 放不下。即使量化到 FP8（70 GB），也需要考虑 KV Cache 的额外显存需求——实际上至少需要两张卡。

更大的模型（如 405B）需要 8 张甚至 16 张卡。如何将计算分布到多张卡上，且保持高效的 GPU 利用率？

两种并行策略各有适用场景：

14.2 张量并行（Tensor Parallelism）

张量并行将每一层的计算切分到多张卡上。核心思想来自 Megatron-LM：利用矩阵乘法的可分性。

列并行（Column Parallel）

源码：vllm/model_executor/layers/linear.py:345

vLLM 的 ColumnParallelLinear 实现了列切分：

python

# vllm/model_executor/layers/linear.py:345-384 (简化)
class ColumnParallelLinear(LinearBase):
    """Linear layer with column parallelism.
    The linear layer Y = XA is parallelized along A's second dimension:
    A = [A_1, ..., A_p].
    """
    def __init__(self, input_size, output_size, ...):
        self.tp_size = get_tensor_model_parallel_world_size()
        # 每张卡只持有 output_size / tp_size 列
        self.output_size_per_partition = output_size // self.tp_size

对于线性层 Y = XW，如果将 W 按列切分为 [W₁ | W₂]：

Y = X × [W₁ | W₂] = [X×W₁ | X×W₂]

GPU 0 计算 X×W₁，GPU 1 计算 X×W₂。输入 X 在两张卡上完全相同（广播），输出 Y 在两张卡上各持有一半。

适用于：QKV 投影（按注意力头切分）、MLP 的 gate/up 投影（按中间维度切分）。

行并行（Row Parallel）

对于线性层 Y = XW，如果将 W 按行切分：

Y = [X₁ X₂] × [W₁; W₂] = X₁×W₁ + X₂×W₂

GPU 0 计算 X₁×W₁，GPU 1 计算 X₂×W₂，然后做 AllReduce（求和）得到最终结果。

适用于：注意力的输出投影、MLP 的 down 投影。

通信开销

张量并行在每个 Transformer 层中引入两次 AllReduce 通信（一次在注意力后，一次在 MLP 后）。在 NVLink 连接的 GPU 之间（如 A100 DGX 的 600 GB/s 双向带宽），这个开销很小。但如果跨机器（通常只有 100 Gbps ≈ 12.5 GB/s），AllReduce 的延迟就会成为显著瓶颈。

经验法则：张量并行只在 NVLink 连接的 GPU 之间使用，不要跨机器。

14.3 流水线并行（Pipeline Parallelism）

流水线并行将模型的不同层放在不同的 GPU 上。数据"流"过各个阶段：

流水线并行只在阶段边界传输一次激活值（而非每层两次 AllReduce），通信量远小于张量并行。因此适合跨机器部署。

但流水线并行有一个固有问题：气泡（Pipeline Bubble）。当 GPU 0 在处理第一批数据时，GPU 1 在空闲；反之亦然。微批次化（将一个批次切成多个微批次）可以减少气泡，但无法完全消除。

vLLM 通过 Ray Compiled DAG 实现跨机器的流水线并行，将多个通信步骤编译为一个固定的执行图，减少了每步的调度开销。

14.4 TP + PP 的组合

实践中，两种并行经常组合使用。例如 8 卡 2 机部署 405B 模型：

机器 A (4 × A100): TP=4, PP stage 0 (Layer 1-40)
机器 B (4 × A100): TP=4, PP stage 1 (Layer 41-80)

机内 TP=4（利用 NVLink 的高带宽），机间 PP=2（跨机通信量最小化）。

通信量对比

理解为什么这样组合，需要比较两种并行的通信模式：

张量并行的通信：每个 Transformer 层有 2 次 AllReduce。假设隐藏维度 d=8192，batch_size=1，每次 AllReduce 传输的数据量 = batch_tokens × d × dtype_size。对于 1 个 Token 的解码步：1 × 8192 × 2 bytes = 16 KB。看起来很小，但 AllReduce 的延迟主要来自启动开销（几微秒），而非数据传输。80 层模型 × 2 次/层 = 160 次 AllReduce/步——NVLink 的低延迟（1-2 μs）可以承受，但跨机器的 RDMA（10-50 μs）会累积到毫秒级。

流水线并行的通信：只在 PP 阶段边界传输一次激活值（batch_tokens × d × dtype_size）。2 个 PP 阶段只有 1 次点对点通信/步，延迟完全可以被 GPU 计算覆盖。

如何选择并行策略

场景	推荐策略	原因
单机 2-8 卡	纯 TP	NVLink 带宽足够，PP 有气泡浪费
2 机 8-16 卡	机内 TP + 机间 PP	最小化跨机通信量
模型能放进单卡	不并行	并行有通信开销
延迟敏感	优先 TP	PP 有流水线气泡增加延迟
吞吐优先	TP + PP	更多卡 = 更多并发

气泡问题的缓解

流水线并行的气泡（bubble）是固有的——当第一个阶段在计算时，最后一个阶段在等待。微批次化（micro-batching）可以减少气泡比例：

微批次数 = 1: 气泡率 ≈ 50%（2 个阶段）
微批次数 = 4: 气泡率 ≈ 20%
微批次数 = 8: 气泡率 ≈ 11%

vLLM 通过连续批处理自然地实现了微批次化——每一步的批次就是一个"微批次"，连续的步骤形成了流水线。

14.5 KV Cache 的分布式传输

当模型采用流水线并行时，一个有趣的问题出现了：KV Cache 需要跟着数据流动吗？

答案是不需要。每个 PP 阶段只负责自己那些层的 KV Cache。Layer 1-20 的 KV Cache 在 GPU 0 上，Layer 21-40 的在 GPU 1 上。数据在阶段间传输的是激活值（隐藏状态），不是 KV Cache。

但在一种新兴的架构——**Prefill-Decode 分离（Disaggregated Serving）**中，KV Cache 确实需要在 GPU 之间传输。这种架构用专门的 GPU 做预填充，完成后将 KV Cache 传输给解码 GPU。vLLM 在 vllm/distributed/kv_transfer/ 中实现了这个机制。

14.6 vLLM 的并行状态管理

源码：vllm/distributed/parallel_state.py

init_distributed_environment()（parallel_state.py:860）是分布式初始化的入口：

python

# vllm/distributed/parallel_state.py:860-895 (简化)
def init_distributed_environment(
    world_size, rank, distributed_init_method="env://",
    local_rank=-1, backend="nccl",
):
    # 如果启用了数据并行(DP)，调整 rank 和 world_size
    if config.parallel_config.data_parallel_size > 1:
        rank = dp_rank * world_size + rank
        world_size = world_size_across_dp

    # 初始化 PyTorch 分布式进程组
    if not torch.distributed.is_initialized():
        torch.distributed.init_process_group(
            backend=backend,          # NCCL for GPU
            init_method=distributed_init_method,
            world_size=world_size,
            rank=rank)

初始化后，每个 Worker 通过 get_tensor_model_parallel_rank() 和 get_pipeline_model_parallel_rank() 获取自己在 TP/PP 组中的位置，决定加载哪些层、切分哪些权重。

vllm/distributed/parallel_state.py 管理所有并行相关的全局状态：

TP 组——同一 PP 阶段内做张量并行的 GPU 组
PP 组——不同 PP 阶段之间的 GPU 组
World 信息——TP rank、PP rank、总 rank 等

Worker 在初始化时根据自己的 rank 加入正确的通信组，之后的 AllReduce、Send/Recv 都在组内完成。

14.7 本章小结

张量并行——切分层内计算，每层 2 次 AllReduce，适合 NVLink 连接的卡
流水线并行——切分层间数据流，只在阶段边界通信，适合跨机器
组合使用——机内 TP + 机间 PP 是大模型部署的标准范式
通信优化——NCCL 原语 + Ray Compiled DAG 减少通信开销

源码导航
并行状态：vllm/distributed/parallel_state.py
通信操作：vllm/distributed/communication_op.py
设备通信器：vllm/distributed/device_communicators/
Ray Executor：vllm/v1/executor/ray_distributed_executor.py

第14章 张量并行与流水线并行 ​

14.1 大模型的困境 ​

14.2 张量并行（Tensor Parallelism） ​

列并行（Column Parallel） ​

行并行（Row Parallel） ​

通信开销 ​

14.3 流水线并行（Pipeline Parallelism） ​

14.4 TP + PP 的组合 ​

通信量对比 ​

如何选择并行策略 ​

气泡问题的缓解 ​

14.5 KV Cache 的分布式传输 ​

14.6 vLLM 的并行状态管理 ​

14.7 本章小结 ​