Skip to content

流水线并行推理优化

流水线并行(Pipeline Parallelism, PP)将模型按层切分到不同 GPU 上,每个 GPU 负责一部分层的计算。与张量并行互补,PP 适用于超大规模模型或跨节点推理场景。

流水线并行推理架构

基本原理

流水线并行将模型的层均匀分配到多个 GPU:

GPU 0: Layer 0-19   → 中间激活发送到 GPU 1
GPU 1: Layer 20-39  → 中间激活发送到 GPU 2
GPU 2: Layer 40-59  → 中间激活发送到 GPU 3
GPU 3: Layer 60-79  → 输出结果

气泡问题与优化

标准流水线存在气泡(Bubble)问题,即部分 GPU 空闲等待:

  • GPipe 策略:将 batch 切分为 micro-batch,填充流水线
  • 1F1B 策略:交替执行前向和反向,减少显存峰值
  • Interleaved 1F1B:每个 GPU 负责不连续的层块,进一步减少气泡
python
# 流水线并行调度示意
class PipelineSchedule:
    def __init__(self, num_stages, num_micro_batches):
        self.num_stages = num_stages
        self.num_micro_batches = num_micro_batches

    def 1f1b_schedule(self):
        """1F1B 调度策略"""
        schedule = []
        # 预热阶段:连续前向
        for i in range(self.num_stages):
            schedule.append(('forward', i, i))

        # 稳态阶段:1前向1反向交替
        for i in range(self.num_micro_batches - self.num_stages):
            schedule.append(('forward', self.num_stages + i, i))
            schedule.append(('backward', i, i))

        # 冷却阶段:连续反向
        for i in range(self.num_stages):
            schedule.append(('backward', self.num_micro_batches - self.num_stages + i, i))

        return schedule

推理中的 PP vs 训练中的 PP

推理场景的流水线并行与训练有所不同:

维度训练 PP推理 PP
反向传播需要不需要
微批次来源梯度累积Continuous Batching
气泡比例10-30%< 5%
通信量大(梯度)小(激活值)

推理场景优势

推理不需要反向传播,因此 PP 的气泡问题远小于训练。结合 Continuous Batching,推理 PP 的 GPU 利用率可达 90% 以上。

配置实践

python
# vLLM 流水线并行配置
from vllm import LLM

llm = LLM(
    model="meta-llama/Llama-2-70b-hf",
    tensor_parallel_size=2,   # 节点内 TP
    pipeline_parallel_size=2, # 跨节点 PP
)

通信开销

PP 的通信是点对点(P2P),数据量为 batch_size × seq_len × hidden_dim。跨节点时 PCIe → NVLink → InfiniBand 的延迟会显著影响性能。建议 PP 仅在单节点内或高带宽互联环境下使用。

相关资源

最近更新