张量切分策略：行切分、列切分、2D切分、混合切分、通信开销分析
好，咱们今天聊聊张量切分。说白了，就是把一个大矩阵拆成几块，分给不同的GPU去算。这玩意儿是模型并行里最基础、也最绕不开的一步。我刚开始接触分布式训练时，觉得切分不就是把矩阵掰开嘛，有啥难的？后来踩了不少坑才发现——切分策略选不对，通信开销能让你怀疑人生。
一、行切分（Row-wise Partition）
行切分，顾名思义，按行把矩阵切开。假设你有一个 [M, N] 的矩阵，切成 P 份，每份就是 [M/P, N]。
什么时候用？ 我个人习惯在数据并行和模型并行混合的场景下用行切分。比如，前向传播时，每个GPU只算自己那几行，最后把结果拼回去。
核心公式：
输入: X [B, M] 权重: W [M, N]
行切分后: W_i [M/P, N]
输出: Y_i = X @ W_i → 形状 [B, N/P] （需要AllReduce合并）嗯，这里要注意：行切分之后，每个GPU拿到的权重是完整的列维度，但行维度变少了。所以计算时，每个GPU只算一部分输出，最后需要做一次AllReduce把结果拼起来。
避坑指南： 我曾经在切分一个超大Embedding层时，没注意行切分后的负载均衡。结果有的GPU分到了高频词，有的分到了低频词，训练速度被最慢的GPU拖死。后来我加了按词频分桶的策略，才把问题解决。
二、列切分（Column-wise Partition）
列切分就是按列切。还是那个 [M, N] 的矩阵，切成 P 份，每份 [M, N/P]。
列切分有个好处——前向传播时，每个GPU算的是完整输出的一部分列，不需要立即通信。但反向传播时，梯度需要做AllReduce。
输入: X [B, M] 权重: W [M, N]
列切分后: W_i [M, N/P]
输出: Y_i = X @ W_i → 形状 [B, N/P] （直接拼接即可，无需通信）你想想看，列切分和行切分最大的区别在哪？行切分需要通信才能得到完整输出，列切分不需要。但列切分在反向传播时，梯度计算需要跨GPU通信。说白了，就是“前向省通信，反向多通信”和“前向多通信，反向省通信”的取舍。
注意： 列切分时，如果输入矩阵的batch size特别大，每个GPU的显存占用会很不均匀。我遇到过一种情况：输入矩阵的某些列特别稀疏，导致有的GPU几乎没算啥，有的GPU算到爆。这时候需要做负载均衡的列重排。
三、2D切分（2D Partition）
2D切分，说白了就是行和列都切。把矩阵切成一个 P_row × P_col 的网格。每个GPU拿到的是一块 [M/P_row, N/P_col] 的子矩阵。
为什么需要2D切分？因为当模型大到一定程度，单一行切分或列切分都不够用。比如一个100B参数的模型，单一行切分可能每个GPU还要装几十GB，显存还是爆。2D切分能把显存占用压到更低。
2D切分的通信模式：
- 行方向：AllReduce（用于梯度同步）
- 列方向：ReduceScatter + AllGather（用于前向/反向计算）
我记得在训练一个千亿参数模型时，我们团队试过纯行切分，结果通信带宽成了瓶颈。后来改成2D切分，虽然通信次数多了，但每次通信的数据量小了，整体吞吐反而提升了30%。
四、混合切分（Hybrid Partition）
混合切分，就是把行切分、列切分、甚至2D切分组合起来用。实际工程中，很少只用一种切分策略。因为模型的不同层，对切分的要求不一样。
举个例子：
- Attention层： 我习惯用列切分。因为Q、K、V的计算可以独立进行，最后拼起来就行。
- FFN层： 用行切分更合适。因为FFN的中间层维度通常很大，行切分能更好地分散计算。
- Embedding层： 用行切分，按词表切。但要注意词频分布。
混合切分的难点在于——不同切分策略之间的转换。比如，Attention层用了列切分，下一层FFN用了行切分，中间就需要做一次重排（Resharding）。这个重排的通信开销，有时候比计算本身还大。
我的经验： 设计混合切分时，尽量让相邻层的切分策略一致，减少重排次数。如果实在避免不了，就把重排操作和计算操作做流水线重叠，别让GPU闲着等数据。
五、通信开销分析
聊了这么多切分策略，核心问题只有一个——通信开销。说白了，切分越细，通信越频繁。但通信带宽是有限的，怎么平衡？
咱们用公式说话：
通信时间 = 数据量 / 带宽 + 延迟 × 通信次数对于行切分：
- 前向：需要AllReduce，通信数据量 = B × N（输出矩阵大小）
- 反向：需要AllReduce，通信数据量 = M × N（梯度矩阵大小）
对于列切分：
- 前向：无需通信（直接拼接）
- 反向：需要AllReduce，通信数据量 = M × N（梯度矩阵大小）
对于2D切分：
- 前向：需要两次通信（ReduceScatter + AllGather）
- 反向：同样两次通信
- 总通信数据量 ≈ 2 × (M × N) / sqrt(P) （P为GPU总数）
| 切分策略 | 前向通信 | 反向通信 | 总通信量 | 适用场景 |
|---|---|---|---|---|
| 行切分 | AllReduce | AllReduce | 2×M×N | 数据并行、小模型 |
| 列切分 | 无 | AllReduce | M×N | Attention层、大batch |
| 2D切分 | RS+AG | RS+AG | 2×M×N/√P | 超大模型、千亿参数 |
| 混合切分 | 视组合而定 | 视组合而定 | 不定 | 复杂模型、异构集群 |
你可能会问：那到底选哪种？我的答案是——没有银弹。我见过有人为了省显存，把模型切得特别碎，结果通信开销占了80%的训练时间，得不偿失。也见过有人为了省通信，用粗粒度的行切分，结果显存爆了。
关键指标： 计算通信比（Compute-to-Communication Ratio）。如果这个比值小于1，说明GPU大部分时间在等数据，切分策略需要调整。我一般要求这个比值至少大于3，否则就考虑合并切分粒度。
六、一张图看懂切分策略
下面我用一张SVG图，把四种切分策略的核心逻辑画出来。你一看就明白。
这张图把四种策略的通信开销和适用场景都标出来了。你注意看右侧的通信开销对比条——行切分最高，2D切分最低。但别忘了，2D切分的实现复杂度也最高。所以实际选型时，我建议你先从简单的行切分开始，等显存不够了再逐步升级。
总结一句话： 切分策略没有绝对的好坏，只有合不合适。关键看你的模型大小、集群带宽、以及你能接受的实现复杂度。我个人习惯是——先跑一个profiling，看看计算通信比，再决定切分策略。别一上来就搞最复杂的，容易把自己绕进去。
好了，张量切分这块就聊到这儿。下一节咱们会深入聊聊具体的通信原语实现，比如AllReduce、ReduceScatter、AllGather这些到底是怎么工作的。到时候我会结合代码讲，保证你听完就能上手。