张量并行推理架构

张量并行（Tensor Parallelism, TP）是将模型权重按维度切分到多个 GPU 上并行计算的策略，是大规模模型推理的核心并行方案。

张量并行推理架构

原理

张量并行的核心思想是将矩阵乘法按列或行切分：

Y = X × W
Y = X × [W1 | W2] = [X×W1 | X×W2]  (列切分)
Y = X × [W1; W2] = X×W1 + X×W2     (行切分，需 AllReduce)

对于 Transformer 模型：

MLP 层：第一层列切分，第二层行切分，仅需一次 AllReduce
Attention 层：QKV 按头切分，Output 行切分，仅需一次 AllReduce

python

# 张量并行 Attention 伪代码
def tp_attention(x, qkv_weights, o_weight, tp_rank, tp_size):
    """张量并行注意力计算"""
    num_heads_per_tp = num_heads // tp_size
    start_head = tp_rank * num_heads_per_tp
    end_head = (tp_rank + 1) * num_heads_per_tp

    # 本地 QKV 计算（无需通信）
    qkv_local = x @ qkv_weights[start_head:end_head]

    # 本地注意力计算
    attn_output = attention(qkv_local)

    # Output 投影（行切分）
    output_local = attn_output @ o_weight[:, start_head:end_head]

    # AllReduce 聚合结果
    output = all_reduce(output_local)

    return output

通信分析

每个 Transformer 层需要 2 次 AllReduce（Attention + MLP 各一次）：

AllReduce 数据量：batch_size × seq_len × hidden_dim × 2 (FP16)
通信次数：2 × num_layers

以 Llama-2-70B 为例，TP=4 时：

每次通信数据量：1 × 1 × 8192 × 2 = 32KB
每层通信次数：2
总通信次数：160（80层 × 2）

通信优化

使用 NCCL 的 Ring AllReduce 算法，TP=4 时每次 AllReduce 的通信量为 3 × (4-1)/4 × data_size。NVLink 带宽 600GB/s 时，通信延迟约 0.05ms，远小于计算时间。

配置实践

python

# vLLM 张量并行配置
from vllm import LLM

llm = LLM(
    model="meta-llama/Llama-2-70b-hf",
    tensor_parallel_size=4,
    gpu_memory_utilization=0.9,
    max_model_len=4096,
)

模型	参数量	推荐 TP	最小 GPU 显存
Llama-2-7B	7B	1	16GB
Llama-2-13B	13B	2	16GB × 2
Llama-2-70B	70B	4	24GB × 4
Llama-2-70B (INT4)	70B	2	24GB × 2

TP 度限制

TP 度必须整除注意力头数。Llama-2-70B 有 64 个头，TP 可选 1/2/4/8/16/32/64。TP=3 是非法配置。

张量并行推理架构 ​

原理 ​

通信分析 ​

配置实践 ​

相关资源 ​

张量并行推理架构

原理

通信分析

配置实践

相关资源