推理集群中的模型分片策略张量并行与流水线并行的通信开销量化对比一、单卡推理 14B 模型直接 OOM分片是唯一的选择但选哪种并行策略部署一个 14B 参数的 LLM 时。FP16 下模型权重约 28GB。加上 KV Cache 和中间激活Activation。单张 24GB 显存的 RTX 4090 直接 OOM。即使是 48GB 的 A6000 也捉襟见肘。模型分片是必然选择。但分片策略的选择不是非黑即白。张量并行Tensor Parallelism, TP和流水线并行Pipeline Parallelism, PP在计算效率和通信开销上有截然不同的特性。选择不当可能导致通信开销吃掉了并行化带来的全部收益。在实际部署中面临的核心问题是给定 N 张 GPU 和模型架构。如何量化两种策略的通信开销。从而做出数据驱动的选型决策。二、张量并行与流水线并行的架构对比两种策略的核心差异在于切分的维度不同。TP 切分的是单个 Transformer 层内部的张量。PP 切分的是层与层之间。graph TD subgraph 张量并行 TP A1[输入] -- B1[GPU 0: W1 分片1] A1 -- B2[GPU 1: W1 分片2] B1 -- C1[All-Reduce(结果)] B2 -- C1 C1 -- D1[GPU 0: W2 分片1] C1 -- D2[GPU 1: W2 分片2] D1 -- E1[All-Reduce(结果)] D2 -- E1 E1 -- F1[输出] end subgraph 流水线并行 PP A2[Micro-Batch 1] -- B3[GPU 0: Layers 0-7] A3[Micro-Batch 2] -- B3 B3 -- C3[GPU 1: Layers 8-15] C3 -- D3[GPU 2: Layers 16-23] D3 -- E3[GPU 3: Layers 24-31] end subgraph 通信模式对比 TP1[TP: 每层都需要 All-Reducebr/通信频率高、数据量小] PP1[PP: 仅层间传激活值br/通信频率低、数据量大] end张量并行的通信模式是高频小量。每个 Transformer 层内部的每个线性层Attention 中的 Q/K/V/O 和 FFN 中的两个 MLP 层都需要一次 All-Reduce 操作。在单次前向传播中一个标准的 Transformer Block 约需要 45 次 All-Reduce。以 32 层模型为例单次前向需 128160 次 All-Reduce。流水线并行的通信模式是低频大量。仅在不同 GPU 上的层间传递完整的激活张量。一个 Token 在 GPU 间传递时。传递的数据量是 hidden_size × batch_size。通信仅发生在每个 Micro-Batch 的边界。单次前向中每 GPU 间仅 1 次通信。关键差异在于通信与计算的重叠能力。TP 的 All-Reduce 可以与后续计算重叠。但重叠程度受限于 All-Reduce 算法的实现Ring vs Tree。PP 的 P2P 通信与计算天然无法重叠。因为下一层的计算必须等待上一层的输出到达。三、通信开销的量化分析公式/// 通信开销的量化计算模型 /// /// 用于在不同模型规模、GPU 数量和网络拓扑下 /// 对比张量并行与流水线并行的通信开销 /// 并策略的通信开销模型 enum ParallelStrategy { /// 张量并行 TensorParallel, /// 流水线并行 PipelineParallel, } /// 模型配置参数 struct ModelConfig { /// 隐藏层维度 hidden_size: usize, /// Transformer 层数 num_layers: usize, /// 注意力头数 num_heads: usize, /// FFN 中间维度通常为 hidden_size × 4 ffn_size: usize, /// 序列长度 seq_len: usize, /// 批处理大小 batch_size: usize, } /// GPU 集群配置 struct ClusterConfig { /// GPU 数量 num_gpus: usize, /// GPU 间带宽GB/s如 NVLink 600GB/s inter_gpu_bw: f64, /// 通信延迟微秒包括启动开销 comm_latency_us: f64, } /// 张量并行的通信开销计算 /// /// 每次 All-Reduce 传输的总量 2 × (g-1)/g × 数据量 /// 其中 g 是 GPU 数量 /// 系数 2 是因为 All-Reduce Reduce-Scatter All-Gather fn calculate_tp_comm_cost(model: ModelConfig, cluster: ClusterConfig) - f64 { let g cluster.num_gpus as f64; // 单层中需要 All-Reduce 的矩阵尺寸按列切分 // QKV 投影: hidden × (3 × hidden/g) let qkv_size model.hidden_size * model.hidden_size * 3; // 输出投影: hidden × hidden let out_size model.hidden_size * model.hidden_size; // FFN 第一层: hidden × (ffn/g) let ffn1_size model.hidden_size * model.ffn_size; // FFN 第二层: ffn × (hidden/g) let ffn2_size model.ffn_size * model.hidden_size; // All-Reduce 通信量 2 × (g-1)/g × 数据量 let allreduce_coeff 2.0 * (g - 1.0) / g; let per_layer_comm_bytes allreduce_coeff * ( qkv_size out_size ffn1_size ffn2_size ) as f64 * 2.0; // FP16 2 bytes // 实际 4 次 All-ReduceQKV、Output、FFN1、FFN2 let per_layer_actual per_layer_comm_bytes / 4.0 * 4.0; let total_comm_bytes per_layer_actual * model.num_layers as f64; // 通信时间 传输量 / 带宽 延迟 × All-Reduce 次数 let transfer_time total_comm_bytes / (cluster.inter_gpu_bw * 1e9); let allreduce_count (model.num_layers * 4) as f64; let latency_time allreduce_count * cluster.comm_latency_us * 1e-6; transfer_time latency_time } /// 流水线并行的通信开销计算 /// /// 每 GPU 间传输一次激活值大小为 batch × seq × hidden × 2 bytes fn calculate_pp_comm_cost(model: ModelConfig, cluster: ClusterConfig) - f64 { let g cluster.num_gpus; // 每 GPU 间的激活值传输量 // 两个方向前向 反向都传输此处仅计算前向 let per_pp_stage_bytes model.batch_size * model.seq_len * model.hidden_size * 2; // FP16 // 总共 g-1 次 GPU 间传输 let total_comm_bytes per_pp_stage_bytes * (g - 1) as f64; let transfer_time total_comm_bytes / (cluster.inter_gpu_bw * 1e9); // PP 的延迟主要是 P2P 发送的启动开销 let latency_time (g - 1) as f64 * cluster.comm_latency_us * 1e-6; transfer_time latency_time } /// 输出对比结果 fn compare_strategies(model: ModelConfig, cluster: ClusterConfig) { let tp_cost calculate_tp_comm_cost(model, cluster); let pp_cost calculate_pp_comm_cost(model, cluster); println!( 通信开销对比 ); println!(张量并行 TP: {:.2} ms, tp_cost * 1000.0); println!(流水线并行 PP: {:.2} ms, pp_cost * 1000.0); println!(比率 TP/PP: {:.2}x, tp_cost / pp_cost); // 示例输出A100 × 4NVLink 600GB/sLlama-14B // 张量并行 TP: 2.34 ms // 流水线并行 PP: 0.18 ms // 比率 TP/PP: 13.0x // // 可以看到 TP 的通信开销是 PP 的 13 倍在这个具体例子中 // 但 TP 的计算效率更高——所有 GPU 同时计算无空闲时间 }量化分析揭示了一个反直觉的结论。虽然 TP 的单次通信量小。但由于频率极高每层 45 次 All-Reduce。总通信开销反而远超 PP。PP 仅传输激活值无论模型多大。每 GPU 间只需传输 batch × seq × hidden 的数据。但这不是说 PP 总优于 TP。TP 的计算效率更高。所有 GPU 在任意时刻都在参与计算。而 PP 存在 Bubble 时间。前面 GPU 处理后续 Micro-Batch 时。后面 GPU 处于空闲。四、分片策略的选择框架与实践建议选择分片策略需要综合考虑模型规模、GPU 拓扑和通信条件。当 GPU 数量较少24 张且使用 NVLink 互联时。TP 是最优选择。NVLink 的 600GB/s 带宽足以支撑 TP 的 All-Reduce 频率。气泡时间成为非问题。当 GPU 数量较多8 张以上且使用 PCIe 或 RoCE 网络时。PP 更优。因为 PCIe 带宽64GB/s和 RoCE200Gbps25GB/s难以支撑 TP 的 All-Reduce 频率。通信会迅速成为瓶颈。在大规模集群16 张以上中。采用混合策略TP PP DP。在同一节点内使用 TPNVLink 互联。节点间使用 PP 和 DP网络互联。这是 DeepSpeed 和 Megatron-LM 等框架推荐的三维并行策略。对于 Batch1 的在线推理场景。TP 的通信开销可能大于计算本身。PP 更加适合。但 PP 的延迟更高Multi-Batch 的流水线填充。需要在延迟和吞吐之间做权衡。五、总结张量并行TP通信频率高每层 45 次 All-Reduce但计算效率高无 Bubble。适合 NVLink 互联的少卡场景。流水线并行PP通信频率低仅层间但有 Bubble 时间。适合网络带宽受限的多机多卡场景。通信开销的量化公式显示 TP 总通信量 2×(g-1)/g × 模型参数量 × 层数。PP 总通信量 (g-1) × batch × seq × hidden。NVLink600GB/s可支撑 TP 通信。PCIe64GB/s和 RoCE25GB/s下 PP 更优。大规模集群推荐 TPPPDP 三维并行。同节点内 TP。节点间 PP/DP。