模型推理的流水线并行不同阶段的延迟不均衡问题与负载均衡策略一、流水线并行的气泡问题在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阶段使每个阶段的总延迟接近。策略二混合切分PPTP。在流水线并行的瓶颈阶段引入张量并行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 jk { 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( fcuda:{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。