模型并行详解：层内模型并行、张量并行、序列并行
模型并行，说白了就是把一个放不下的模型拆开，放到多张卡上一起跑。我刚开始接触分布式训练时，总觉得数据并行就够了，直到有一次要训练一个百亿参数的Transformer模型，单卡显存直接爆掉。嗯，从那以后我才真正开始研究模型并行。
今天咱们聊聊三种主流的模型并行方式：层内模型并行、张量并行、序列并行。它们解决的是不同维度的问题，但目标一致——让大模型能跑起来，而且跑得快。
一、层内模型并行：最朴素的切分思路
层内模型并行，也叫算子内并行。它的思路很简单：把一个计算层内部的参数切分开，放到不同设备上计算。
举个例子，一个全连接层的权重矩阵是 W，形状是 [d_in, d_out]。我们可以把它按列切成两块：W1 和 W2，分别放到 GPU0 和 GPU1 上。输入 x 同时发给两块卡，各自算完后再把结果拼接起来。
核心要点：层内模型并行不改变网络层数，只改变单层内部的参数分布。
我在项目中遇到过一个问题：层内并行虽然能降低单卡显存，但通信开销很大。每次前向和反向都要做 all-reduce 或 all-gather，如果跨机网络带宽不够，反而比单卡还慢。
我的建议：层内并行适合单机多卡场景，跨机时慎用。如果必须跨机，优先考虑后面要讲的张量并行。
二、张量并行：更精细的切分策略
张量并行是层内并行的升级版。它把张量运算本身拆开，让每张卡只算一部分，最后再合并结果。
以 Megatron-LM 的 1D 张量并行为例，它把 Transformer 中的 QKV 投影矩阵按列切分，把输出投影矩阵按行切分。这样做的好处是：前向计算中只需要一次 all-reduce，通信效率更高。
# 伪代码：1D张量并行前向
def forward(x, W_qkv, W_out):
# 每张卡持有 W_qkv 的一部分
qkv_local = x @ W_qkv_local # 本地计算
# all-reduce 合并结果
qkv = all_reduce(qkv_local)
# 后续计算...
你想想看，为什么 Megatron 选择这种切分方式？因为 Transformer 的注意力计算中，QKV 投影是独立的，切分后每张卡算自己的部分，最后 all-reduce 一下就行。通信量只有一次，效率很高。
注意：张量并行对通信带宽要求很高。我建议在单机内使用 NVLink 或 NVSwitch 连接时采用，跨机时通信延迟会明显增加。
除了 1D 张量并行，还有 2D 和 3D 版本。2D 把矩阵按行列同时切分，3D 再加一个深度维度。维度越高，通信模式越复杂，但理论上能支持更大的模型。
| 并行方式 | 切分维度 | 通信次数 | 适用场景 |
|---|---|---|---|
| 1D 张量并行 | 行或列 | 1次 all-reduce | 单机多卡，NVLink |
| 2D 张量并行 | 行和列 | 2次 all-reduce | 更大模型，多机 |
| 3D 张量并行 | 行、列、深度 | 3次 all-reduce | 超大规模训练 |
三、序列并行：处理长序列的新思路
序列并行是最近几年才火起来的。它解决的是序列长度带来的显存问题。
Transformer 的显存占用和序列长度是平方关系。序列长度翻倍，注意力计算的显存占用翻四倍。我训练一个 128K 长度的模型时，单卡显存直接飙到 80GB，根本放不下。
序列并行的思路是：把序列切分成多个片段，每张卡处理一个片段。但这里有个问题——注意力计算需要全局信息，片段之间怎么交互？
目前主流做法有两种：
- Ring Attention：把序列切成块，每张卡算自己的块，然后通过环形通信交换 key-value 对，逐步完成全局注意力计算。
- Flash Attention + 序列并行：利用 Flash Attention 的分块计算特性，把不同块分配到不同设备上，最后合并结果。
关键点：序列并行不减少计算量，但能显著降低单卡显存。它和张量并行、数据并行可以组合使用。
我曾经踩过一个坑：序列并行时，如果序列长度不是设备数的整数倍，会导致负载不均。有些卡算得快，有些卡算得慢，整体效率被拖垮。后来我加了一个动态负载均衡策略，才解决这个问题。
四、三种并行的对比与选择
这三种并行方式不是互斥的，实际训练中经常组合使用。我个人的习惯是：
- 模型小于 10B 参数：只用数据并行
- 10B ~ 100B 参数：数据并行 + 张量并行
- 100B 以上：数据并行 + 张量并行 + 序列并行 + 流水线并行
说白了，没有银弹。每种并行方式都有它的适用场景和代价。你需要根据模型大小、硬件配置、网络带宽来权衡。
避坑指南：我曾经为了追求极致的显存节省，同时用了 3D 张量并行和序列并行，结果通信开销占了 70% 的训练时间。后来我意识到，显存省下来了，但训练速度慢得无法接受。所以，并行策略的选择一定要考虑通信开销。
五、知识体系总览
下面这张图展示了三种模型并行的核心逻辑和关系：
这张图把三种并行的切分方式、通信模式和适用场景都列出来了。你可以把它当作一个快速参考。
最后说一句：模型并行的选择没有标准答案。我建议你先从简单的开始，比如数据并行 + 张量并行，等遇到显存瓶颈再逐步加入序列并行。别一开始就搞太复杂的组合，调试起来会让人崩溃。
想要更多资料，联系可可