模型并行基础：张量并行（Tensor Parallelism）的原理与实现

张量并行，说白了就是把一个巨大的矩阵运算，切分成几块，让不同的GPU各自算一块。我刚开始接触这个概念时，总觉得这不就是矩阵分块乘法吗？后来真正上手才发现，里面的门道比想象中多得多。

为什么需要张量并行？

你想想看，现在的大模型动辄几百亿参数。单张GPU的显存，就算H100也才80GB。一个175B的GPT-3，光参数就要350GB。模型并行就派上用场了。

模型并行有两种主流方式：

• 流水线并行（Pipeline Parallelism）：按层切分，每张卡负责几层
• 张量并行（Tensor Parallelism）：按矩阵维度切分，每张卡负责一部分计算

我个人习惯把张量并行看作是「更细粒度的模型并行」。它解决的是单层计算量太大的问题。举个例子，一个隐藏层维度为12288的FFN层，单卡算起来显存压力巨大，但切成4份后，每张卡只算3072维，压力就小多了。

张量并行的两种切分方式

这里我重点讲两种最常见的切分策略。我在项目中遇到过不少同学搞混它们，结果通信开销比计算还大，得不偿失。

1. 按行切分（Row-wise Parallelism）

假设我们要计算 Y = X × W，其中X是输入，W是权重矩阵。

按行切分就是把W按行切成两块：

W = [W1]    →   GPU0: Y1 = X × W1
    [W2]        GPU1: Y2 = X × W2
                AllReduce: Y = Y1 + Y2

这里有个关键点：每张卡都需要完整的输入X。所以输入需要广播到所有卡上。计算完成后，需要做一次AllReduce把结果合并。

2. 按列切分（Column-wise Parallelism）

按列切分则是把W按列切开：

W = [W1 | W2]  →   GPU0: Y1 = X × W1
                    GPU1: Y2 = X × W2
                    拼接: Y = [Y1 | Y2]

每张卡只需要计算一部分输出维度。最后通过AllGather把结果拼起来就行。

嗯，这里要注意：按列切分不需要AllReduce，只需要AllGather。通信量比行切分小一些。我一般优先考虑按列切分，除非模型结构不允许。

Transformer中的张量并行实现

我们拿Transformer的一个Transformer Block来举例。它包含两个主要部分：多头注意力（MHA）和前馈网络（FFN）。

先看MHA层的张量并行：

# 伪代码：4卡张量并行下的MHA
class TensorParallelMHA:
    def __init__(self, world_size=4):
        self.world_size = world_size
        # QKV权重按列切分
        self.q_proj = ColumnParallelLinear(dim, dim // world_size)
        self.k_proj = ColumnParallelLinear(dim, dim // world_size)
        self.v_proj = ColumnParallelLinear(dim, dim // world_size)
        # 输出投影按行切分
        self.out_proj = RowParallelLinear(dim, dim)
    
    def forward(self, x):
        # 每张卡计算部分QKV
        q = self.q_proj(x)  # 输出维度: dim/world_size
        k = self.k_proj(x)
        v = self.v_proj(x)
        
        # 本地计算注意力
        attn_output = self_attention(q, k, v)
        
        # 输出投影需要AllReduce
        output = self.out_proj(attn_output)
        return output

FFN层的实现类似：

class TensorParallelFFN:
    def __init__(self, world_size=4):
        # 第一个线性层：按列切分
        self.gate_proj = ColumnParallelLinear(dim, 4*dim // world_size)
        self.up_proj = ColumnParallelLinear(dim, 4*dim // world_size)
        # 第二个线性层：按行切分
        self.down_proj = RowParallelLinear(4*dim, dim)
    
    def forward(self, x):
        gate = self.gate_proj(x)
        up = self.up_proj(x)
        hidden = silu(gate) * up
        output = self.down_proj(hidden)
        return output

通信开销分析

张量并行不是免费的午餐。它引入了额外的通信开销。我整理了一个对比表：

切分方式	前向通信	反向通信	通信量
按行切分	AllReduce	AllReduce	2 × 输出大小
按列切分	AllGather	ReduceScatter	1 × 输出大小

你看，按列切分的通信量只有按行切分的一半。所以能按列切分的地方，我尽量按列切分。

实际部署中的注意事项

我在部署千亿模型时，总结了几条经验：

1. 通信拓扑很重要：张量并行要求卡间通信带宽极高。最好用NVLink连接的卡做张量并行，跨节点用流水线并行。
2. 切分粒度要合适：不是切得越细越好。切太细了，通信开销会吃掉计算加速。我一般控制在4-8卡之间。
3. 注意显存碎片：张量并行后，每张卡的显存使用不均衡。我曾经遇到过某张卡显存爆了，其他卡还有富余的情况。需要仔细规划切分策略。

张量并行整体架构

下面这张图展示了张量并行在Transformer Block中的完整流程：

这张图展示了一个完整的张量并行流程。输入先广播到4张卡上，每张卡独立计算部分QKV和注意力，然后通过AllGather合并，再经过FFN层，最后AllReduce得到最终输出。

性能调优建议

最后分享几个我在实际调优中总结的点：

• 计算通信重叠：把通信操作和计算操作重叠起来。比如在AllReduce的同时，让GPU继续算下一层。这需要仔细设计流水线。
• 选择合适的通信后端：NCCL是首选，但如果你用的是AMD的卡，得用RCCL。我踩过这个坑，换了后端后性能提升明显。
• 注意负载均衡：确保每张卡的计算量差不多。如果某张卡算得快，其他卡算得慢，快的卡就得空等。