大模型推理中的张量并行：详解 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 利用率低	小规模测试，PCIe 连接较差环境
Gemm 重叠	利用矩阵乘法特性，减少等待时间	需要算子融合支持，调试复杂	高吞吐推理，支持 Tensor Core 的显卡
请求间重叠	提升多并发下的整体吞吐量	调度逻辑复杂，显存占用增加	高并发服务，多租户场景
请求内重叠	降低单个请求的延迟	实现难度最高，对 Attention 结构改动大	低延迟敏感型应用，长序列处理

总结

张量并行通信计算重叠是提升大模型推理性能的关键技术之一。随着硬件互联带宽的提升（如 NVLink）以及软件栈的成熟，通信计算重叠将成为标准配置。开发者在选择推理引擎或进行自研优化时，应重点关注其对不同重叠模式的支持程度，并根据实际部署环境的网络拓扑和负载特征进行权衡。

未来，随着 MoE 架构和大模型规模的进一步扩大，通信瓶颈将更加突出，细粒度的通信计算重叠优化将是持续演进的方向。

大模型推理中的张量并行：详解 4 种通信计算重叠模式