📚 课程目录第4章张量并行原理

4、张量并行原理：张量并行切分方式、通信模式（AllReduce等）及实现细节

4.1 核心原理：从“单卡放不下”到“按维度切分”

当单个GPU的显存无法容纳大模型的权重或中间激活时，除了流水线并行按“层”纵向切分，另一种更细粒度的方案是张量并行（Tensor Parallelism, TP）。其核心思想是：将某一层的权重矩阵按行或按列切分成多个分片，分别存放在不同GPU上，让多张卡协同完成同一层的计算。与流水线并行不同，张量并行是层内并行，每张卡只持有权重的一部分，但通过集合通信（Collective Communication）在计算过程中交换中间结果，从而在数学上等价于完整权重的计算。

以Transformer中的线性层 Y = X · W 为例，设输入 X 的形状为 [batch, seq_len, hidden]，权重 W 的形状为 [hidden, output]。张量并行有两种基本切分方式：

按列切分（Column Parallel）：将 W 沿输出维度切成 TP_size 块，每块形状为 [hidden, output/TP_size]。每张卡独立计算 Y_local = X · W_local，得到输出的一部分。此时不需要通信，因为输入X是完整的，各卡只算输出的一部分。
按行切分（Row Parallel）：将 W 沿输入维度切成 TP_size 块，每块形状为 [hidden/TP_size, output]。此时输入 X 也必须按列切分，每张卡持有 X_local，计算 Y_local = X_local · W_local，得到的是部分和，需要经过一次 AllReduce 将各卡的部分和相加，才能得到完整的 Y。

在实际的Transformer层中，通常将列并行用于第一个线性层（如QKV投影），将行并行用于第二个线性层（如注意力输出投影或FFN的第二层），这样可以在两层之间自然地衔接，减少不必要的通信。

  关键认知：张量并行不是“把模型复制到多卡”，而是“把同一层的权重拆开，让多卡合算一层”。通信开销与切分维度、TP大小、隐藏层尺寸直接相关。

4.2 通信模式：AllReduce 与 AllGather 的取舍

张量并行中最核心的通信原语是 AllReduce。它完成“多卡各持部分结果 → 求和（或求平均等）→ 每卡都得到完整结果”的操作。在NCCL（NVIDIA Collective Communications Library）中，AllReduce 有多种实现算法，常见的有：

Ring AllReduce：将GPU组成一个环，分两步完成——Reduce-Scatter（各卡得到部分和）和 AllGather（各卡广播得到完整结果）。通信量为 2 × (TP_size - 1) / TP_size × 数据量，带宽利用率高，适合大规模TP。
Tree AllReduce：基于树形结构做规约和广播，延迟较低，适合小数据量或延迟敏感场景。

除了 AllReduce，张量并行中还会用到 AllGather（各卡把自己的分片广播给所有卡，最终每卡持有完整数据）和 Reduce-Scatter（各卡的部分结果按目标卡规约，最终每卡持有不同的部分和）。在TensorRT-LLM的MLP（多层感知机）实现中，通常采用“列并行 + 行并行”的组合，中间只需要一次 AllReduce，而不是两次，从而将通信量减半。

工程经验：在TensorRT-LLM中，对于FFN（前馈网络）层，通常将第一个线性层按列切分（无通信），第二个线性层按行切分（需要AllReduce）。这样整个FFN只产生一次AllReduce通信，而不是两次。

4.3 方法流程：TensorRT-LLM 中的张量并行实现细节

在TensorRT-LLM中，张量并行的实现涉及权重加载时的切分和推理时的通信插入两个阶段。以下是关键流程：

权重切分：模型加载时，根据 tensor_parallel_size（TP大小）和 tensor_parallel_rank（当前GPU的rank），将每个张量按预设的切分模式（行/列）切分。TensorRT-LLM通过 sharding_dim 属性指定切分维度。
构建Engine：TensorRT在构建引擎时，将通信算子（如AllReduce）作为网络中的节点插入。每个GPU只保存自己的权重分片，但计算图是完整的。
运行时通信：推理时，每张卡执行自己的计算子图，遇到通信节点时调用NCCL进行集合通信。TensorRT-LLM使用 ncclComm 管理通信组。

一个典型的张量并行配置示例如下：

# 构建TensorRT-LLM引擎时指定TP大小
trtllm-build --model_dir ./llama-7b-hf \
             --tensor_parallel_size 4 \
             --output_dir ./llama-7b-tp4

# 运行时指定rank
mpirun -n 4 python run.py \
       --engine_dir ./llama-7b-tp4 \
       --tensor_parallel_rank $RANK

在代码层面，TensorRT-LLM的 Attention 层中，QKV投影采用列并行，输出投影采用行并行。对于 MLP 层，gate和up投影采用列并行，down投影采用行并行。这种设计使得每个Transformer块中只出现两次AllReduce（一次在Attention输出，一次在MLP输出），通信开销被最小化。

4.4 案例/场景：LLaMA-7B 在 4 卡上的张量并行

以LLaMA-7B为例，其隐藏层维度为4096，FFN中间维度为11008。若使用 TP=4：

层/操作	权重形状	切分方式	每卡分片形状	通信
QKV投影	[4096, 12288]	按列切分	[4096, 3072]	无
Attention输出投影	[4096, 4096]	按行切分	[1024, 4096]	AllReduce
gate/up投影	[4096, 11008]	按列切分	[4096, 2752]	无
down投影	[11008, 4096]	按行切分	[2752, 4096]	AllReduce

从表中可以看出，每个Transformer层仅产生2次AllReduce通信。每次AllReduce的数据量为 batch × seq_len × 4096 × 4字节（FP32）或 × 2字节（FP16）。对于 batch=1, seq_len=2048 的输入，单次AllReduce数据量约为 1×2048×4096×2 = 16MB，在NVLink（约300GB/s）上耗时约0.05ms，而计算耗时通常在几毫秒量级，通信占比可控。

4.5 误区/复盘：张量并行常见陷阱

误区一：认为TP越大越好。TP的通信量随TP大小近似线性增长（Ring AllReduce的通信量为 2×(TP-1)/TP×数据量），当TP超过8时，通信开销可能超过计算收益，尤其在单节点内NVLink带宽有限的情况下。通常TP=2或4是性价比最高的选择，更大规模应结合流水线并行。

误区二：忽略切分维度的正确性。对于非对称权重（如Embedding和LM Head），切分方式不同。Embedding通常按行切分（词表维度），LM Head按列切分，且需要保证词表维度对齐。若切分错误，会导致推理结果错误且难以排查。

误区三：忘记在AllReduce后同步。在自定义实现中，如果AllReduce没有正确同步，后续计算可能读到未规约完成的中间值。TensorRT-LLM的图优化会自动插入同步点，但手工修改图时需格外小心。

复盘建议：在调试张量并行模型时，先用 TP=1 跑通基线，再逐步增加TP大小，对比每层输出是否一致。可以使用 tensorrt_llm 提供的 check_tensor_parallel 工具或自行打印中间张量的 allclose 结果来验证切分和通信的正确性。