模型推理的流水线并行：不同阶段的延迟不均衡问题与负载均衡策略

一、流水线并行的气泡问题

在Transformer大模型推理中，当模型参数量超过单个GPU显存容量时，必须将模型切分到多个设备上。流水线并行（Pipeline Parallelism, PP）按层切分模型——GPU-0负责第1-8层，GPU-1负责第9-16层，以此类推——每张GPU仅持有模型的连续片段。推理请求按批次流经各张GPU，形成处理流水线。

理论上，一个有K个阶段的流水线应当实现K倍的吞吐提升。但现实中，流水线的吞吐受限于最慢的阶段（木桶原理），且阶段切换之间存在"气泡"（Bubble）——下游GPU等待上游GPU输出时的空闲时间。在推理场景中，气泡问题比训练更为严重，因为推理无法通过增加micro-batch数来摊平气泡。

二、延迟不均衡的成因分析

流水线中各阶段的延迟差异来自三个层面：

计算量差异。Transformer各层的计算量理论上相同，但模型切分并不总是均匀的——当总层数不能被GPU数整除时，部分GPU被分配更多层。此外，embedding层和lm_head层的计算量通常小于中间层，导致首尾GPU利用不足。

显存带宽差异。即使计算量相同，层在GPU上的实际执行时间还取决于KV Cache的大小。在推理中，较早的层处理短序列（未经过位置编码的累计），较晚的层处理长序列——KV Cache的读写开销随序列长度递增。

设备异构性。在混合GPU集群中（如A100和V100混合），不同设备的计算能力差异直接映射为阶段延迟差异。

import torch import time from typing import List, Tuple, Dict def profile_pipeline_latency( model_segments: List[torch.nn.Module], devices: List[torch.device], sample_input: torch.Tensor, num_warmup: int = 10, num_measure: int = 100 ) -> Dict[int, float]: """Profile流水线各阶段的延迟。 在目标设备上实际执行每个模型片段，测量平均延迟。 用于诊断流水线中的瓶颈阶段。 Args: model_segments: 切分后的模型片段列表 devices: 每个片段对应的目标设备 sample_input: 标准化的测试输入 num_warmup: 预热迭代次数 num_measure: 测量迭代次数 Returns: {stage_id: avg_latency_ms} 的阶段延迟字典 """ stage_latencies = {} for stage_id, (segment, device) in enumerate( zip(model_segments, devices) ): segment = segment.to(device) segment.eval() # 准备该阶段的输入 stage_input = sample_input.to(device) # 预热GPU和CUDA缓存 with torch.no_grad(): for _ in range(num_warmup): _ = segment(stage_input) torch.cuda.synchronize(device) # 精确计时 latencies = [] with torch.no_grad(): for _ in range(num_measure): torch.cuda.synchronize(device) t0 = time.perf_counter() _ = segment(stage_input) torch.cuda.synchronize(device) latencies.append( (time.perf_counter() - t0) * 1000 # ms ) stage_latencies[stage_id] = sum(latencies) / len(latencies) # 释放该设备的显存 del segment torch.cuda.empty_cache() return stage_latencies

三、负载均衡策略：从静态切分到动态调度

策略一：计算量加权切分。最简单的改进：在切分模型时，不是按层数均分，而是基于profile结果按实际延迟分配。将延迟更高的层分配给更少的GPU阶段，使每个阶段的总延迟接近。

策略二：混合切分（PP+TP）。在流水线并行的瓶颈阶段引入张量并行（Tensor Parallelism, TP），即将该阶段的计算进一步拆分到多个GPU上并行执行。这以额外的通信为代价，降低了瓶颈阶段的延迟。

策略三：多流水线交错。维护多个并行的流水线实例，当一条流水线中某阶段空闲时，交错插入另一流水线中的请求。这在vLLM和TensorRT-LLM中有对应实现。

def balance_pipeline_by_latency( layer_latencies: List[float], # 每层的延迟（ms） num_stages: int ) -> List[List[int]]: """基于动态规划的延迟均衡切分。 目标：将层分配到num_stages个阶段，最小化各阶段 延迟的最大值与最小值的差距（minimize max-min）。 使用动态规划求解最优分配： dp[i][k] = 前i层分配到k个阶段的最小最大延迟 Args: layer_latencies: 每层的独立延迟（ms） num_stages: 流水线阶段数（GPU数） Returns: stage_assignments: 每个阶段的层索引列表 """ n = len(layer_latencies) # 前缀和加速区间求和 prefix_sum = [0] * (n + 1) for i in range(n): prefix_sum[i + 1] = prefix_sum[i] + layer_latencies[i] # dp[i][k] = min over j<k { max(dp[j][k-1], sum(layer_latencies[j:i])) } dp = [[float('inf')] * (num_stages + 1) for _ in range(n + 1)] split = [[0] * (num_stages + 1) for _ in range(n + 1)] dp[0][0] = 0 for i in range(1, n + 1): for k in range(1, min(i, num_stages) + 1): for j in range(k - 1, i): stage_latency = prefix_sum[i] - prefix_sum[j] max_lat = max(dp[j][k - 1], stage_latency) if max_lat < dp[i][k]: dp[i][k] = max_lat split[i][k] = j # 回溯分配方案 assignments = [] i = n for k in range(num_stages, 0, -1): j = split[i][k] assignments.append(list(range(j, i))) i = j return list(reversed(assignments))

四、通信与计算的重叠优化

流水线中一个关键的优化是：在等待上游GPU输出时，当前GPU可以提前执行不依赖上游结果的操作。例如，KV Cache的预分配、output tensor的显存空间预留等工作可以与数据传输重叠。

在实现层面，(1) 使用torch.cuda.Stream创建多个CUDA流，将通信放在一个流、计算放在另一个流，实现真正异步重叠；(2) 使用send/recv而非基于all_reduce的集合通信——流水线中数据是单向的P2P传输，使用点对点通信指令更高效。

def create_overlapped_pipeline_stage( segment: torch.nn.Module, prev_device: int, next_device: int ): """创建通信-计算重叠的流水线阶段。 使用双CUDA流：compute_stream处理层计算， comm_stream处理跨设备数据传输。 两个流可以并行执行，实现计算和通信的重叠。 Args: segment: 当前阶段的模型片段 prev_device: 上游设备编号 next_device: 下游设备编号 """ compute_stream = torch.cuda.Stream() comm_stream = torch.cuda.Stream() def process_batch(hidden_states: torch.Tensor) -> torch.Tensor: """处理一个batch的数据，重叠通信与计算。 流程： 1. (comm_stream) 接收上一阶段的输出 → 拷贝到当前设备 2. (compute_stream) 执行当前阶段的计算 3. (comm_stream) 将当前阶段的输出发送到下游设备 """ # Step 1: 在通信流中接收数据 with torch.cuda.stream(comm_stream): # 确保数据在当前设备上 current_input = hidden_states.to( f"cuda:{torch.cuda.current_device()}" ) # Step 2: 在计算流中处理（与后续通信可重叠） with torch.cuda.stream(compute_stream): output = segment(current_input) # 等待计算完成 torch.cuda.current_stream().wait_stream(compute_stream) return output return process_batch

五、总结

流水线并行的负载均衡本质上是一个调度优化问题。在实践中，三条经验法则值得关注：(1) 在切分模型前先用小batch profile各层在各GPU上的实际延迟——纸面计算量不能替代实际测量，KV Cache和显存带宽的影响可能颠覆纯计算量模型的预测；(2) 瓶颈阶段引入张量并行是成本最低的补救措施——增加两张GPU做TP通常比重新设计整个流水线分配要简单；(3) 多流水线交错策略在batch数量足够大时效果显著，但每个batch的端到端延迟不会缩短——它提升的是吞吐（throughput）而非单个请求的延迟（latency）。