大模型推理中的张量并行：详解 4 种通信计算重叠模式
背景
张量并行（Tensor Parallelism, TP）目前已经是大模型推理的一个必备技术。然而，张量并行存在一个显著的缺点，即通信开销。当推理采用 PCIe 类卡互联时，该缺点更加明显。
针对通信开销的缺点，训练框架已经有了通信计算重叠优化（Communication-Computation Overlap），而目前部分开源的推理引擎如 vLLM 和 SGLang 早期版本并未完全实现该功能。近期开源的大模型推理引擎 ZhiLight 支持张量并行通信计算重叠。预计在 2025 年，张量并行通信计算重叠将会是所有主流开源框架的必备功能。
本文结合当前最新的论文与工程实践，介绍张量并行通信计算重叠的几种做法。
张量并行的几种实现
1. 朴素版张量并行
标准的 Transformer 张量并行结构中，每次 Transformer 前向需要进行 2 次 AllReduce 操作。这会导致模型前向执行 AllReduce 的时候，计算的 GPU 处于空闲状态，造成资源浪费。
2. Gemm 版本通信计算重叠
当我们说到张量并行计算通信重叠，一个最直观的实现是分布式 Gemm + AllReduce 的重叠。目前 TransformerEngine、PyTorch (TorchTitan) 和字节 Flux 都是采取类似的实现。
在分布式 Gemm + AllReduce 中，A @ B 的计算过程如下：
- A 是列切分，B 是行切分。
- A0 @ B0 得到 C00 部分，A @ B1 得到 C01 部分。
- 蓝色的 C00/C11 和黄色的 C01/C11 进行 ReduceScatter 得到最终的 C0/C1。
在重叠版本中，原来的 A 按照列切分，计算的时候再按照行分块计算。分为两个 step：
- Step 0: A00 @ B0 得到 C00，A01 @ B 得到 C01。
- Step 1: A10 @ B0 得到 C10，A11 @ B 得到 C11。 此时可以同时进行 Step 0 计算结果的规约（Reduce）。Flux 的思想与此类似，但在优化方面做得更细致。
3. 请求间通信计算重叠
这是张量并行通信与计算重叠的另一种实现方式（参考 Liger: Interleaving Intra- and Inter-Operator Parallelism for Distributed Large Model Inference）。该方法的特点包括：
- 会有多个请求，不同的请求会有不同的 Stream。
- 执行请求 1 的计算的时候，请求 2 正在进行通信操作，反之亦然。
类似的这种做法还有 Nanoflow。从理论上讲，这种方法不需要重写一个计算通信的 kernel，但是整体调度实现会很复杂。
4. 请求内通信计算重叠
第 3 种通信计算重叠方式如 ISO: Overlap of Computation and Communication within Sequence For LLM Inference 所述。该方法看起来与前几种均不一样，核心在于对分布式 Attention 实现的深入理解。
以单张卡不同流的图示为例，将这张图扩展到多卡场景：
- Rank 0 和 Rank 1 分别处理不同的数据分片。
- 一个 Transformer 对于单卡而言，通信和计算是重叠的。
- Attention 采用了分块 Attention 实现。
具体流程为：每张卡在序列维度分块执行。分块 0 执行的时候没有通信，分块 1 执行的时候，执行分块 0 的通信。MLP 的计算通信重叠通常与 Gemm 版本类似，重点在于 Attention 如何实现通信计算重叠。
方案对比与选型建议
在实际工程中，选择哪种重叠策略取决于具体的硬件环境和业务需求：
| 方案 | 优点 | 缺点 | 适用场景 |
|---|---|---|---|
| 朴素版 | 实现简单，兼容性好 | 通信阻塞严重，GPU 利用率低 |
