流水线并行推理优化
流水线并行(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 仅在单节点内或高带宽互联环境下使用。