张量并行（Tensor Parallelism）推理部署实战：从 Megatron-LM 切分原理到 vLLM 多卡配置与性能调优

GPU张量并行推理部署架构示意图

当大语言模型的参数量从 7B 一路攀升到 70B 甚至 405B 时，单张 GPU 的显存早已无法容纳完整模型。以 Llama-3-70B 为例，仅 FP16 权重就需要约 140GB 显存，远超单张 A100-80G 的容量上限。此时，将模型拆分到多张 GPU 上并行推理成为唯一可行的方案，而张量并行（Tensor Parallelism, TP）正是其中最核心、也是工程实践中使用频率最高的并行策略。

本文将从 Megatron-LM 的切分原理出发，深入讲解张量并行的数学基础与工程实现，然后以 vLLM 为主要框架，覆盖从配置部署到性能调优的完整流程，帮助你真正跑通多卡推理的生产级部署。

一、张量并行的核心思想：把矩阵乘法拆到多张卡上

张量并行的本质是将模型中的矩阵运算（GEMM）沿某个维度切分，分配到不同 GPU 上各自计算一部分，再通过通信原语合并结果。与数据并行（Data Parallelism）每张卡持有完整模型不同，TP 下每张卡只持有模型权重的一个切片——这对显存受限的推理场景至关重要。

Transformer 模型中需要并行化的核心组件有两个：多层感知机（MLP）和多头注意力（Multi-Head Attention）。Megatron-LM 论文提出了优雅的切分方案，分别对应列并行（Column-Parallel）和行并行（Row-Parallel）线性层。

1.1 列并行线性层（Column-Parallel Linear）

对于线性层

Y = XA

，其中输入矩阵

维度为

(m, k)

，权重矩阵

维度为

(k, n)

。列并行将

沿输出维度

切分为

份：


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
# 列并行切分示意

# A: (k, n) -> [A_1, A_2, ..., A_N], 每个 A_i: (k, n/N)

# 每张 GPU 计算 Y_i = X * A_i, 输出 Y_i: (m, n/N)

# 无需通信，各卡独立计算



import torch

import torch.nn as nn



class ColumnParallelLinear(nn.Module):

    def __init__(self, in_features, out_features, world_size):

        super().__init__()

        assert out_features % world_size == 0

        self.out_per_partition = out_features // world_size

        self.weight = nn.Parameter(torch.empty(self.out_per_partition, in_features))

        self.bias = nn.Parameter(torch.empty(self.out_per_partition))



    def forward(self, x):

        # x: (m, k) -- 每张卡都有完整的输入

        # weight: (n/N, k) -- 每张卡只有一部分权重

        return torch.nn.functional.linear(x, self.weight, self.bias)

关键点在于：输入 X 在每张卡上是完整的，每张卡用自己的权重切片独立计算，输出

Y_i

的维度是

(m, n/N)

。这个阶段完全不需要通信。

1.2 行并行线性层（Row-Parallel Linear）

行并行将权重矩阵

沿输入维度

切分，每张卡持有一行切片

A_i

，维度为

(k/N, n)

。同时输入

也沿

维切分：


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
# 行并行切分示意

# A: (k, n) -> 沿 k 维切分为 [A_1, ..., A_N], 每个 A_i: (k/N, n)

# X: (m, k) -> 沿 k 维切分为 [X_1, ..., X_N], 每个 X_i: (m, k/N)

# 每张 GPU 计算 Y_i = X_i * A_i, 输出 Y_i: (m, n)

# 需要 All-Reduce 将所有 Y_i 求和得到最终 Y



class RowParallelLinear(nn.Module):

    def __init__(self, in_features, out_features, world_size):

        super().__init__()

        assert in_features % world_size == 0

        self.in_per_partition = in_features // world_size

        self.weight = nn.Parameter(torch.empty(out_features, self.in_per_partition))

        self.bias = nn.Parameter(torch.empty(out_features))



    def forward(self, x_partition):

        # x_partition: (m, k/N) -- 每张卡只有一部分输入

        # 计算部分结果

        partial_y = torch.nn.functional.linear(x_partition, self.weight)

        # 需要 All-Reduce 聚合所有卡的结果

        return partial_y

行并行计算后，每张卡得到的是部分和，必须通过 All-Reduce 操作将所有部分和相加才能得到最终结果。这就是 TP 中通信开销的主要来源。

张量并行列并行与行并行矩阵切分示意图

二、Megatron-LM 的 MLP 切分方案：列并行 + 行并行 = 零冗余通信

Transformer 的 MLP 模块由两个线性层组成：

1	FFN(x) = Dropout(GELU(x * A) * B)

，其中

维度为

(k, 4k)

，

维度为

(4k, k)

（隐含 4 倍扩展比）。Megatron-LM 的精妙之处在于：第一个线性层用列并行，第二个用行并行，中间不需要任何通信。


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
class MegatronMLP(nn.Module):

    def __init__(self, hidden_size, world_size):

        super().__init__()

        # 第一个线性层: 列并行, 每张卡输出 (4k/N) 维

        self.fc1 = ColumnParallelLinear(hidden_size, 4 * hidden_size, world_size)

        # 第二个线性层: 行并行, 每张卡输入 (4k/N) 维

        self.fc2 = RowParallelLinear(4 * hidden_size, hidden_size, world_size)

        self.gelu = nn.GELU()

        self.dropout = nn.Dropout(0.0)  # 推理时 dropout=0



    def forward(self, x):

        # x: (m, k) -- 完整输入, 每张卡都有

        intermediate = self.gelu(self.fc1(x))  # (m, 4k/N) -- 列并行输出

        output = self.fc2(intermediate)        # (m, k) -- 行并行部分和

        # 只需一次 All-Reduce 即可得到最终结果

        return all_reduce(output)

整个 MLP 前向过程中，只在最后做一次 All-Reduce。中间的 GELU 激活函数直接在列并行的输出上局部计算，不需要任何通信。这是 Megatron-LM 方案高效的核心原因。

2.1 Multi-Head Attention 的并行化

注意力层的并行化更加直观——直接将不同的注意力头分配到不同 GPU 上。如果模型有 32 个头、TP size = 4，则每张卡负责 8 个头的 QKV 计算和 Attention 计算。QKV 投影使用列并行（输出维度按头数切分），输出投影使用行并行（输入维度按头数切分），与 MLP 的模式完全一致。

组件	第一个线性层	第二个线性层	通信次数
MLP	列并行 (k → 4k/N)	行并行 (4k/N → k)	1 次 All-Reduce
Attention (QKV)	列并行 (k → k/N 按头切)	行并行 (k/N → k)	1 次 All-Reduce
每层总通信	2 次 All-Reduce（MLP + Attention 各一次）

三、推理场景下的张量并行 vs 训练场景：关键差异

虽然 TP 的切分原理在训练和推理中是相同的，但工程实现上有几个重要差异，理解这些差异是做好推理部署的前提。

差异一：没有梯度通信。训练时每层需要做两次 All-Reduce（前向 + 反向），推理只需要前向的一次 All-Reduce，通信量减半。

差异二：权重固定不变。推理时模型权重是静态的，可以在启动时一次性切分好，不需要在运行时动态同步。这意味着 TP 的通信开销只来自激活值的 All-Reduce，与权重规模无关。

差异三：Decode 阶段的通信效率问题。在自回归解码阶段，每步只生成一个 token，batch_size 通常较小（甚至为 1），此时每次 All-Reduce 传输的数据量极小（几十 KB），但通信的固定延迟（latency）不变。这导致 小 batch decode 场景下 TP 的通信开销占比极高，是性能瓶颈的主要来源。

场景	Prefill (大batch)	Decode (小batch)
计算量	高（GEMM 密集）	低（GEMV 为主）
通信量	大（MB 级）	小（KB 级）
瓶颈	计算受限	通信延迟受限
TP 效率	高（计算/通信比好）	低（通信延迟占比高）

这也是为什么 PD 分离架构会将 Prefill 和 Decode 部署在不同集群上——Decode 集群可以选用较小的 TP size，甚至用单卡部署多个副本，以降低通信延迟的影响。

四、vLLM 张量并行部署实战

vLLM 是当前最流行的开源推理框架之一，对张量并行有原生支持。下面从零开始演示如何部署一个 TP=4 的 70B 模型推理服务。

4.1 基本启动命令


1
2
3
4
5
6
# 使用 4 张 GPU 部署 Llama-3-70B (TP=4)

vllm serve meta-llama/Meta-Llama-3-70B \

  --tensor-parallel-size 4 \

  --gpu-memory-utilization 0.90 \

  --max-model-len 8192 \

  --port 8000

关键参数说明：

1
--tensor-parallel-size 4

：设置 TP size 为 4，模型权重将沿 4 张 GPU 切分。vLLM 内部会自动按 Megatron 方案对 MLP 和 Attention 层进行列并行/行并行切分。
1
--gpu-memory-utilization 0.90

：控制每张 GPU 用于推理的显存比例（含 KV Cache）。70B 模型在 TP=4 下每卡约需 35GB 权重，留出空间给 KV Cache。
1
--max-model-len 8192

：最大上下文长度，直接影响 KV Cache 的预分配量。

4.2 TP Size 的选择策略

选择 TP size 不是越大越好。核心原则是：刚好能装下模型即可。原因有二：

第一，TP size 越大，每张卡的权重越少，但 All-Reduce 通信次数不变（每层 2 次），通信延迟开销随 TP size 线性增长（Ring All-Reduce 的延迟与参与节点数成正比）。

第二，更大的 TP size 意味着更低的每卡计算密度。Decode 阶段每卡只做极小的 GEMV，计算资源利用率很低。


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# TP size 选择参考表（FP16 推理）

# 模型参数量 -> 最少 GPU 数 -> 推荐 TP size



tp_recommendations = {

    "7B":  {"min_gpus": 1, "recommended_tp": 1, "vram_per_gpu": "~14GB"},

    "13B": {"min_gpus": 1, "recommended_tp": 2, "vram_per_gpu": "~13GB"},

    "34B": {"min_gpus": 1, "recommended_tp": 2, "vram_per_gpu": "~34GB"},

    "70B": {"min_gpus": 1, "recommended_tp": 4, "vram_per_gpu": "~35GB"},  # A100-80G

    # 70B on A100-40G: 需要 TP=4 (每卡 35GB > 40GB 留余量给KV Cache)

    "70B_40G": {"min_gpus": 4, "recommended_tp": 4, "vram_per_gpu": "~35GB"},

    # 405B: 需要 TP=8 (每卡 ~101GB, 需要 H100-80G + 量化)

    "405B": {"min_gpus": 8, "recommended_tp": 8, "vram_per_gpu": "~101GB"},

}



for model, config in tp_recommendations.items():

    print(f"{model:12s} -> TP={config['recommended_tp']}, {config['vram_per_gpu']}/GPU")

4.3 模型权重转换：从 HuggingFace 到 Megatron 格式

vLLM 内部使用 Megatron 格式的权重布局。大多数情况下，vLLM 能自动从 HuggingFace 格式转换，但对于某些自定义模型或需要极致启动速度的场景，可以预先转换：


1
2
3
4
5
6
7
8
# vLLM 自动转换（首次加载稍慢，后续有缓存）

vllm serve meta-llama/Meta-Llama-3-70B --tensor-parallel-size 4



# 对于 TensorRT-LLM, 需要手动转换权重

python3 tensorrt_llm/examples/llama/convert_checkpoint.py \

  --model_dir /path/to/hf-llama-70b \

  --output_dir /path/to/trt-llama-70b-tp4 \

  --tp_size 4

多GPU服务器张量并行推理部署架构图

五、性能调优：通信优化与吞吐量提升

5.1 NCCL 通信优化

张量并行的 All-Reduce 依赖 NCCL 库。合理的 NCCL 配置可以显著降低通信延迟：


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
# 1. 确保使用 NVLink/NVSwitch 互联（而非 PCIe）

# 检查 GPU 互联拓扑

nvidia-smi nvlink -s

nvidia-smi topo -m



# 2. 设置 NCCL 环境变量优化通信

export NCCL_DEBUG=INFO          # 首次运行时查看通信拓扑和带宽

export NCCL_IB_DISABLE=1        # 单机多卡不需要 InfiniBand

export NCCL_NET_GDR_LEVEL=PHB   # GPU Direct RDMA 级别

export NCCL_SOCKET_IFNAME=lo    # 单机使用 loopback

export NCCL_BUFFSIZE=8388608    # 增大 NCCL buffer (8MB)



# 3. 使用 CUDA Graph 减少内核启动开销

# vLLM 已内置支持, 确保启用

vllm serve meta-llama/Meta-Llama-3-70B \

  --tensor-parallel-size 4 \

  --enforce-eager  # 注意: 设为 False(默认) 才启用 CUDA Graph

5.2 Pipeline Parallelism 配合 TP 使用

当模型大到单机 GPU 数量不够时（如 405B 模型需要 8+ 张卡），需要跨机部署。此时纯 TP 不够高效（跨机 NVLink 不存在，All-Reduce 走网络延迟极高），需要引入流水线并行（Pipeline Parallelism, PP）：


1
2
3
4
5
6
7
8
9
10
11
12
13
# 跨 2 台机器, 每台 4 张 GPU, TP=4 PP=2

# 机器0 (rank 0-3): 模型的 layer 0-31

# 机器1 (rank 4-7): 模型的 layer 32-63



# 机器0

NCCL_DEBUG=WARN vllm serve meta-llama/Meta-Llama-3-70B \

  --tensor-parallel-size 4 \

  --pipeline-parallel-size 2 \

  --distributed-executor-backend ray \

  --ray-workers-use-gpu



# 机器1 需要启动 ray worker

ray start --address='<机器0_IP>:6379' --gpu

PP 将模型按层切分到不同机器上，机器间只传输激活值（而非完整的 All-Reduce），通信量远小于跨机 TP。组合策略通常是：机内用 TP（利用 NVLink），机间用 PP（减少网络通信）。

5.3 Benchmark 与性能验证

部署完成后，务必进行基准测试以验证 TP 配置是否合理：


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# 使用 vLLM 自带的 benchmark 脚本

# 测试吞吐量 (tokens/s)

python3 benchmarks/benchmark_throughput.py \

  --backend vllm \

  --model meta-llama/Meta-Llama-3-70B \

  --tensor-parallel-size 4 \

  --input-len 1024 \

  --output-len 512



# 测试延迟 (ms/token, 单请求)

python3 benchmarks/benchmark_latency.py \

  --backend vllm \

  --model meta-llama/Meta-Llama-3-70B \

  --tensor-parallel-size 4 \

  --batch-size 1 \

  --output-len 128

预期性能参考（A100-80G x4, TP=4, Llama-3-70B, FP16）：

指标	典型值	说明
吞吐量 (大batch)	~2000-3000 tokens/s	batch_size=32, input=1024, output=512
延迟 (单请求)	~40-60 ms/token	batch_size=1, decode 阶段
Prefill 延迟	~200-500 ms	input=1024 tokens
显存利用率	~90%	含 KV Cache

六、常见问题与排障指南

6.1 OOM（显存不足）

最常见的问题。排查步骤：


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# 1. 检查每张卡的实际显存占用

nvidia-smi --query-gpu=memory.used,memory.total --format=csv



# 2. 降低 gpu-memory-utilization (但会减少 KV Cache 空间)

vllm serve ... --gpu-memory-utilization 0.85



# 3. 减小 max-model-len

vllm serve ... --max-model-len 4096



# 4. 增大 TP size (如果有多余 GPU)

vllm serve ... --tensor-parallel-size 8



# 5. 使用量化 (AWQ/GPTQ INT4 可将显存减半)

vllm serve TheBloke/Llama-3-70B-AWQ \

  --tensor-parallel-size 2 \

  --quantization awq

6.2 NCCL 通信超时


1
2
3
4
5
6
7
# 设置更长的 NCCL 超时

export NCCL_TIMEOUT=1800  # 30 分钟 (用于大模型权重加载)



# 或者通过 vLLM 参数

vllm serve ... \

  --tensor-parallel-size 4 \

  --distributed-timeout 1800

6.3 TP size 必须能整除注意力头数

这是一个容易被忽略的硬性约束。如果模型有 32 个注意力头，TP size 可以是 1、2、4、8、16、32，但不能是 3、5、6 等。部署前务必检查模型的

1	num_attention_heads

配置：


1
2
3
4
5
6
7
# 检查模型的注意力头数

from transformers import AutoConfig

config = AutoConfig.from_pretrained("meta-llama/Meta-Llama-3-70B")

print(f"num_attention_heads: {config.num_attention_heads}")

print(f"valid TP sizes: {[h for h in range(1, config.num_attention_heads + 1) 

                        if config.num_attention_heads % h == 0]}")

# 输出示例: valid TP sizes: [1, 2, 4, 8, 16, 32, 64]

总结：张量并行推理部署的核心要点

张量并行是大模型推理部署的基础能力。回顾本文的关键要点：

切分原理：Megatron-LM 的列并行 + 行并行方案，每层只需一次 All-Reduce，是最优的 TP 切分策略。
TP size 选择：刚好装下模型即可，过大的 TP size 会增加通信延迟而降低效率。Decode 阶段尤其敏感。
框架差异：vLLM、SGLang、TensorRT-LLM 都支持 TP，但配置方式和优化深度不同。vLLM 最易用，TensorRT-LLM 性能最优。
通信优化：确保 NVLink 互联、合理设置 NCCL 环境变量、启用 CUDA Graph，可以显著降低通信开销。
多机扩展：机内用 TP、机间用 PP，是跨机部署的标准范式。Ray 是 vLLM 多机部署的推荐后端。
排障：OOM 是最常见问题，可通过降低 KV Cache、增大 TP、使用量化来缓解。注意 TP size 必须整除注意力头数。

掌握这些原理和实操技巧，你就能在各种 GPU 配置下稳定地部署大模型推理服务。如果你对 PD 分离架构或更细粒度的通信优化感兴趣，推荐阅读本站的PD 分离部署实战和NCCL 通信原语详解，进一步深入分布式推理的工程实践。