4、模型并行策略：张量并行、流水线并行、序列并行与1F1B调度

各位同学，今天我们来聊聊模型并行。说实话，单卡训练大模型的时代早就过去了。我刚开始做分布式训练那会儿，一个7B的模型硬塞进一张A100，显存直接爆掉。后来我学乖了——模型太大？那就把它拆开，让多张卡一起扛。

模型并行，说白了就是把一个模型切成几块，分给不同的GPU去算。但怎么切、切完怎么调度，这里面门道很深。我个人习惯把模型并行分成三种：张量并行、流水线并行、序列并行。咱们一个一个来看。

4.1 张量并行（Tensor Parallelism）

张量并行，我更喜欢叫它“层内切分”。什么意思呢？就是把一个层的计算，比如矩阵乘法，拆成多个小矩阵，分给不同的GPU去算。

举个例子，一个线性层 Y = XW，其中W的维度是 [d_in, d_out]。我们可以把W按列切成两块：W = [W1, W2]，每块维度是 [d_in, d_out/2]。然后GPU0算 Y1 = XW1，GPU1算 Y2 = XW2。最后把Y1和Y2拼起来，就得到了完整的Y。

这样做的好处很明显：每张卡只存一半的权重，显存压力直接减半。但代价呢？每次前向和反向传播，都需要一次all-reduce通信，把各卡的梯度汇总。我在项目中遇到过，如果通信带宽不够，张量并行的收益会被通信开销吃掉。

代码实现上，PyTorch的 TensorParallel 模块已经封装好了。但如果你自己手写，核心就是 torch.distributed.all_reduce 和 torch.chunk 的配合。

# 伪代码示例：张量并行的前向传播
def forward(self, x):
    # 每张卡只计算自己那部分权重
    local_output = torch.matmul(x, self.weight_chunk)
    # 所有卡的结果做all-reduce求和
    torch.distributed.all_reduce(local_output)
    return local_output

4.2 流水线并行（Pipeline Parallelism）

流水线并行，就是“层间切分”。把模型按层切成几段，每张卡负责一段。比如一个12层的Transformer，切成4段，每张卡负责3层。

但这里有个坑：如果按最朴素的方式，一张卡算完才交给下一张卡，那大部分时间GPU都在空等。这就是所谓的“气泡”（bubble）问题。我曾经在项目里踩过这个坑，显存利用率低得可怜。

怎么解决？引入微批次（micro-batch）。把一个大batch切成多个小batch，让流水线“流”起来。GPU0算完第一个微批次，立刻传给GPU1，同时开始算第二个微批次。这样流水线就填满了。

4.3 1F1B调度

说到流水线调度，就不得不提1F1B（One Forward One Backward）。这是目前最主流的调度策略，也是我个人的首选。

1F1B的核心思想很简单：每个GPU交替执行一次前向传播和一次反向传播。这样做的好处是，前向和反向的计算可以重叠，减少GPU空闲时间。

具体来说，每个GPU维护一个队列。前向传播的结果先存着，等反向传播的梯度回来，再一起算。这样每个GPU的显存占用是稳定的，不会出现前向堆积导致显存爆炸的情况。

我记得有一次，团队里有人用了传统的“先全部前向再全部反向”的调度，结果显存直接翻倍。换成1F1B后，显存占用降了40%，训练速度还快了20%。

4.4 序列并行（Sequence Parallelism）

序列并行，这是最近比较火的一个方向。它解决的是长序列训练的问题。你想想看，一个序列长度达到8192甚至16384时，注意力机制的显存占用是O(n²)的，单卡根本扛不住。

序列并行的思路是：把序列长度切成多段，每张卡负责一段。但这里有个难点——注意力计算需要全局信息，不能简单切分。所以序列并行通常和张量并行配合使用。

具体做法是：在注意力层，把序列维度切分到多张卡上，每张卡计算局部注意力，然后通过通信交换key-value对，得到全局注意力结果。这样每张卡只需要处理一部分序列，显存占用从O(n²)降到了O((n/p)²)，其中p是卡数。

我在实际项目中测试过，序列长度16384时，用8卡序列并行，显存占用从80GB降到了12GB。效果非常显著。

最后，我建议大家在选择并行策略时，先看看你的瓶颈在哪。显存不够？优先考虑张量并行或序列并行。计算效率低？试试流水线并行加1F1B调度。没有银弹，只有最适合你场景的方案。

好了，这一章的内容就到这里。记住，模型并行不是银弹，但用好了能解决大模型训练中80%的显存问题。下一章我们会聊聊混合并行——怎么把数据并行、张量并行、流水线并行组合起来，发挥最大威力。