张量并行（Tensor Parallelism，TP）大概是分布式大模型训练里最“烧脑”的一环，很多人第一次接触它的感觉是：概念听了一堆，AllReduce也认识，但真要自己画一个Transformer的结构图，把每个矩阵怎么切、每张卡上留什么、什么时候通信讲明白，就懵了。这篇文章就是干这个的。
我会用大量图的方式，而不是死磕公式，从最简单的矩阵乘法开始拆，逐步讲完一个Transformer block里每个Linear层怎么配置TP，再把通信量怎么算、跟数据并行怎么搭配、踩坑点在哪一次说清楚。不管你是在读源码，还是想自己动手改并行策略，这篇文章都能给你一个能直接上手的“脑内模型”。
1. TP到底在切什么——从一次矩阵乘法说起
1.1 图1：一个 Linear 层拆成两个 GPU 去算
先看最基本的数学事实。一个线性层做的事情是 Y = X·W，X 是输入序列，形状是 [seq_len, H]，W 是权重矩阵，形状是 [H, N]，Y 是输出，形状是 [seq_len, N]。
如果只有一张卡，这个矩阵乘法算完就结束了。现在有两张卡，怎么分？最直觉的做法是把 W 切成两块，让每张卡算一半。
图1画的就是这个场景：左边是原始计算，一个大矩阵 W 完整放在一张卡里；右边是切分后的样子，W 被从中间一刀劈开，变成 W1 和 W2 两块，分别放在 GPU0 和 GPU1 上，然后 GPU0 算出 X·W1，GPU1 算出 X·W2。
这里有一个关键细节：无论切成几块，每一步计算都是“全量X 乘以 分块W”，所以每张卡都需要拿到完整的 X。也就是权重是分着存的，但输入是每个进卡都要有一份的。当你把 W 切成 2 份，单卡显存里 Linear 参数的占用就直接除以 2 了——这是 TP 最核心的价值。
但是，问题立刻来了：GPU0 输出的是 Y 的前 N/2 列，GPU1 输出的是后 N/2 列，这俩只是“竖着拼在一起”才能还原出完整 Y 的，拼起来需要一次通信。什么样的通信？后面会讲，这里先记住“拼接”这个动作。
1.2 图2：竖着切和横着切，通信方式完全不一样
图2展示的是另一种切法。刚才我们竖着切了 W——也就是把输出维度 N 切开，每一卡算输出的“一部分列”。现在我们把 W 横着切，也就是把输入维度 H 切开。
这时候情况变了。为了能让 GPU0 算出完整的 Y 列，它需要 W 的上半部分和 X 的左半部分；GPU1 需要 W 的下半部分和 X 的右半部分。两边各自算出一个“部分和”，比如 GPU0 算出 Y_a，GPU1 算出 Y_b，真正的 Y = Y_a + Y_b。
这个“相加”的动作，在分布式里就是 AllReduce。两个 GPU 各拿一个部分值，AllReduce 做完之后，每张卡上都有完整、相同的 Y。
所以图2要传达的核心信息是：
- 竖切 W（列切分）：每张卡从完整输入算部分输出，要拼起来 → 用 AllGather。
- 横切 W（行切分）：每张卡从部分输入算完整输出，要加起来 → 用 AllReduce。
这俩在工程实现里有一个不对称的优劣关系。列切分的通信量看着更小（因为只要拼起来），但用了列切分之后每张卡手里只有一部分输出，下一个模块如果要用完整输出，还得再通信；行切分则直接把结果“广播”到每张卡上了，虽然 AllReduce 通信量稍大，但下游计算可以直接开始。
一个 Linear 层怎么选切分方式，不取决于这个层本身，取决于它的输出接下来怎么被使用。这是理解 TP 布局的关键钥匙，先记住这句话。
1.3 为什么 TP 能省显存：不只是权重那一份
很多人算显存时容易漏掉梯度、优化器状态和激活值。TP 省显存不是只省了权重，而是连同反向传播时的梯度也一起省了。W 被切到两张卡上之后，反向时每张卡只需要计算自己那一块 W 的梯度，梯度矩阵天然也只有一半大小。
优化器状态是跟着参数走的。以 Adam 为例，一份参数至少对应两份优化器状态（momentum 和 variance），参数量是模型参数量的好几倍。TP 把参数切碎，优化器状态也随之分散到多张卡上。
激活值这块则要分开说。因为 TP 是“每个采样batch里、每条序列上的同一个Transformer层被拆到多卡上”，激活值也会被拆分到对应的维度上。比如 QKV 投影的输出被列切分之后，激活值也分到各卡。但有些激活值是没法切的，比如经过 LayerNorm 前的完整输入X，每卡都需要一份，这个显存省不掉。
实际训练里，一个百亿参数模型用 TP=8 跑起来，单卡能比纯数据并行明显多塞不少 batch size，正因为上述四项都被分散了。不过 TP 的代价是通信频繁，后面专门讲。
2. Transformer 里的 TP 布局——QKV、O、MLP 各自怎么切
2.1 图3：一个 Transformer block 的双 TP 布局总览
先把整个 Transformer block 画出来：输入 X → LayerNorm → QKV 投影 → 多头注意力 → O 投影 → 残差加 → LayerNorm → MLP 扩展层 → 激活函数 → MLP 压缩层 → 残差加 → 输出。
图3标出了这个 block 里每个线性层的 TP 切分方式。如果你打开 Megatron 风格的实现，会看到非常规律的配置：
- QKV 投影：列切分，输出被拆到多卡，每卡负责一部分注意力头。
- O 投影：行切分，输入来自各卡的部分头输出，输出是完整结果，内部做一次 AllReduce。
- MLP 扩展层：行切分，输入是完整 X，输出是“每卡一份的隐藏层分片”。
- MLP 压缩层：列切分，输入是隐藏层分片，输出是完整 Y，内部做一次 AllReduce。
也就是说，一个 Transformer block 里，前一个层负责把数据“切出去”，后一个层负责把数据“收回来”。QKV 和 O 是一对，MLP 扩展层和压缩层是一对。
这也解释了为什么图3里每个 block 只有两次 AllReduce——一次在 Attention 的 O 投影，一次在 MLP 压缩层。这两次通信是“必须合流”的地方，其他位置数据都是分片状态，各自算各自的。
2.2 图4：Attention 的 TP 细节——为什么 QKV 能直接各算各的
图4把注意力部分放大。假设模型有 H 个注意力头，TP 是 2，每张卡拿 H/2 个头。QKV 投影用列切分，权重被切到两卡上，每卡只负责自己那部分头的 Q、K、V。
关键点：注意力头之间本来是完全独立的——head i 的 Q 只跟 head i 的 K、V 做 attention，不会跟 head j 交互。所以每个卡只需拿着自己的那部分 QKV 权重，算自己的头，完全不需要跨卡通信。这个“无通信”区间从 QKV 投影开始，一直持续到所有头计算完毕、拼接出完整的注意力输出为止。
拼接完 attention 输出之后，O 投影要把它映射回完整的 hidden size。O 投影用行切分，每张卡只持有一半的“注意力输出分片”和一半的行权重，各自算出完整输出的“部分和”，然后 AllReduce 一次得到完整结果。
这里有一个初学者容易绕晕的点：既然注意力头是独立的，为什么 QKV 投影不用行切分，这样每张卡不就能持有一份完整的输出吗？答案是没必要，也不能。如果用行切分，输入 X 也要被切，但 X 是上层 O 投影 AllReduce 出来的完整结果，你已经把它广播到每张卡了，再强行切一次就是白通信。QKV 用列切分，正好让每个卡只保留自己的头所需要的输出维度，省了拼接动作，让数据直接进入注意力计算。
2.3 图5：MLP 的 TP 细节——先展宽再还原的经典两段式
MLP 的结构是 扩展层（把 H 映射到 4H）→ 激活 → 压缩层（把 4H 映回 H）。图5画的就是两段式的 TP。
扩展层用行切分。输入是完整 X，权重被横切成两半，每卡只有 4H/2 的“部分权重”，算出来的输出是 4H 维度的一半。这两半其实正好对应着隐藏层的两个分片，彼此不干扰，所以不需要立刻通信，直接进激活函数就行。
压缩层用列切分。它的输入是“扩展层输出的那个 4H 分片”，所以每卡只需要拿一半的行权重，就能把自己那一半隐藏层映射成一个完整 H 维度的部分结果。最后 AllReduce 合并，得到完整 Y。
把扩展层和压缩层合在一起来看，通信模式特别整齐：扩展层把数据“拆散”，压缩层把数据“收回”。MLP 的唯一一次 AllReduce 发生在压缩层之后，也就是在残差相加之前。残差相加需要完整的 Y，所以那次 AllReduce 不可避免。
一个常见的优化是：把残差相加和 AllReduce 融合，或者把 AllReduce 跟后一个 LayerNorm 的统计量计算融合。很多框架做了这个融合，目的就是减少内核启动次数，省掉通信等待的延迟。
3. 代码级实现——手写一套 TP Linear 层
3.1 基础工具：进程组初始化与权重切分
在看图理解了布局之后，动手写代码会透彻很多。下面用 PyTorch 写一个极简版本，忽略框架层面的封装，重点展示“数据是怎么流动的”。
首先初始化进程组，每个进程就是一张卡：
import torch import torch.distributed as dist def setup(rank, world_size): dist.init_process_group( backend="nccl", init_method="env://", rank=rank, world_size=world_size, ) # 假设已经在4卡环境中，rank从0到3，world_size=4然后定义一个工具函数，把权重切出一个“分片”。这里用最简单的方式：每张卡直接按自己的 rank 取权重块，实际工程里会有更精细的保存/加载逻辑，但核心思想一样：
def get_weight_chunk(world_size, rank, full_weight, dim): # full_weight: 完整权重矩阵 chunk_size = full_weight.size(dim) // world_size return full_weight.narrow(dim, rank * chunk_size, chunk_size).clone()把一次矩阵乘法的输入 X 广播给所有卡，就可以开始列切分/行切分的实现了。
3.2 ColumnParallelLinear 实现
列切分就是“竖着切 W”。每张卡拿着完整的 X 和 W 的一个竖条，算自己那部分输出，最后 AllGather 拼接完整输出：
class ColumnParallelLinear(torch.nn.Module): def __init__(self, in_features, out_features, world_size, rank): super().__init__() self.world_size = world_size self.rank = rank # 完整权重只存在于“模拟”层面，这里直接创建全量再切 full_weight = torch.randn(out_features, in_features) / (in_features ** 0.5) self.weight = torch.nn.Parameter( get_weight_chunk(world_size, rank, full_weight, dim=0) ) def forward(self, x): # x: [seq_len, in_features]，所有卡持有完整 x local_out = torch.matmul(x, self.weight.t()) # [seq_len, out_features / world_size] # AllGather 拼接列 gathered = [torch.empty_like(local_out) for _ in range(self.world_size)] dist.all_gather(gathered, local_out) return torch.cat(gathered, dim=-1)注意两个细节：
self.weight的形状是[out_features / world_size, in_features]，所以x和它的转置相乘后得到局部输出。all_gather会把所有节点上的local_out收集起来，顺序按 rank 排列，拼接后就是完整输出。
这个实现直接完成了 1.1 图1中“竖切”的全部流程。
3.3 RowParallelLinear 实现
行切分是“横着切 W”。每张卡拿着部分 X 和部分 W，做一次局部矩阵乘，然后 AllReduce 加总：
class RowParallelLinear(torch.nn.Module): def __init__(self, in_features, out_features, world_size, rank): super().__init__() self.world_size = world_size self.rank = rank full_weight = torch.randn(out_features, in_features) / (in_features ** 0.5) self.weight = torch.nn.Parameter( get_weight_chunk(world_size, rank, full_weight, dim=1) ) def forward(self, x_partial): # x_partial: [seq_len, in_features / world_size] local_out = torch.matmul(x_partial, self.weight.t()) # [seq_len, out_features] dist.all_reduce(local_out) return local_out这个实现的关键是，x_partial必须是前一个模块切出来的“那一份”，而不是完整 X。这里weight形状是[out_features, in_features / world_size]。
Column 和 Row 两个类加起来，正好对应 2.1 图3里的四层布局：QKV 投影用 Column，O 投影用 Row，MLP 扩展层用 Row，MLP 压缩层用 Column。
3.4 一个最小可跑的 MLP 验证
把上面两个层拼成一个极小的两层 MLP，就可以验证 TP 的正确性：
class TPMLP(torch.nn.Module): def __init__(self, hidden_dim, world_size, rank): super().__init__() # 扩展层：行切分。输入完整X，输出分片 self.fc1 = RowParallelLinear(hidden_dim, hidden_dim * 4, world_size, rank) # 压缩层：列切分。输入分片，输出完整，内部AllReduce self.fc2 = ColumnParallelLinear(hidden_dim * 4, hidden_dim, world_size, rank) def forward(self, x): h = self.fc1(x) h = torch.relu(h) return self.fc2(h)这段代码在 4 卡上跑一次前向，每个卡的最终输出是完全一致的——因为它们内部做了 AllReduce。你可以把自己的权重初始化为固定值，对比它和单卡版本的输出，误差应该在浮点范围内接近于零。
这个最小验证非常值得跑一下。我自己最初学 TP 的时候，就是因为把 Row 和 Column 的切分方向搞反了，出来的结果是 NaNs。如果你跑出来数值对不上，优先检查权重切分方向和输入 X 是不是“该完整就完整、该分片就分片”。
4. TP、DP、PP 怎么选——并行策略对比与显存核算
4.1 三种并行策略本质对比
图6画的是数据并行（DP）、张量并行（TP）、流水线并行（PP）三种策略下，模型和数据的分布方式。
- 数据并行：每张卡都有完整模型，但喂不同的数据，算完同步梯度。显存冗余大，但通信量相对可控。
- 张量并行：模型权重被切碎分散到多卡，所有卡共同处理同一份数据，靠密集通信实时合并结果。
- 流水线并行：模型按层切成多个阶段，每个阶段放在不同的卡上，数据像一个流水线一样逐段流过。
三者不是互斥的，现代大模型训练基本都是三维组合：先用 TP 把单层算力拼起来，再用 PP 把层切成多段，外部再用 DP 扩展吞吐。
那“什么时候用 TP”呢？答案取决于你的单卡显存是否能够放下一个 Transformer 层。如果放不下，DP 再多卡也没用，因为每张卡都需要完整的一层权重；PP 当然可以帮你分段，但同一层内部的矩阵乘法还是太大了。TP 是唯一能对“单个算子级别”进行拆分的方案。所以模型规模一大，TP 就是绕不开的。
我自己的经验是：当模型参数量超过单卡显存能承载的权重+激活值后，先用 TP 把每层拆到能装下为止，如果 TP 世界规模过大导致通信浪费，再用 PP 把层切开。这个顺序很重要，很多人搞反了，上来就用 PP，结果发现单层照样爆显存。
4.2 显存核算实例：一个 7B 模型怎么放上 8 张卡
假设有一个 7B 参数的模型，参数精度 BF16，每参数 2 字节，那么权重至少 14GB。加上梯度 14GB、Adam 优化器状态 28GB，训练总计至少需要 56GB。单卡如果只有 40GB，根本跑不了。
用 TP=2，每卡参数降到 7GB，梯度 7GB，优化器 14GB，一共 28GB，单卡 40GB 能塞下。但这还没算激活值。激活值跟 batch size、序列长度强相关，通常需要几 GB 到十几 GB。如果再叠加 PP=2，把部分激活值进一步分摊，单卡压力会更小。
这里有个工程建议：准备训练配置时，先把“参数+梯度+优化器”算清楚，这个基本是硬性的；激活值可以通过减小 batch size 或使用激活检查点来调节。TP 主要是为了压缩前三者，激活值问题主要靠序列并行和激活检查点去缓解。
4.3 TP 规模的边界在哪里
TP 不能无限加大。每切一刀，权重虽然小了，但通信量上升了，而且矩阵乘法里的“无效计算”比例也在升高。图7里我画了一条曲线：横轴是 TP 规模，纵轴是“算力利用率”。典型的曲线上，TP=2 到 TP=4 时利用率还不错，TP=8 之后开始明显下滑，超过 TP=16 时通信开销往往会吃掉大部分并行收益。
同时，TP 要求多卡之间通信带宽高、延迟低。NVMe/PCIe 级别的互联跑 TP 会很痛苦。通常 TP 只应该在单机内跨卡使用，通过 NVLink 连接；跨节点的 TP 基本不划算，跨节点时应该用 PP 或 DP 去扩展。
很多框架的默认并行配置也是这么做的：单机内先 TP，多机之间 PP，机与机之间 DP。这套组合拳是有原因的，因为它把“最高频的通信”（TP 的 AllReduce）留在带宽最高的 NVLink 域内，把“低频通信”放到跨机网络上。
5. 通信量压榨——从 AllReduce 到分解式低通信模式
5.1 图7：AllReduce 的内部结构
标准的 AllReduce 语义是“所有节点输入的张量按元素相加，再把结果广播回所有节点”。一个最直接实现方式叫 Ring-AllReduce，流程是先 ReduceScatter 再 AllGather。
想象一个环上有 4 张卡，每张卡持有一部分数据块。ReduceScatter 阶段，数据块沿着环转一圈，每个节点把接收到的数据和自己的数据相加，最后每个节点持有的是“总和后的一个分块”。AllGather 阶段，各节点把那个分块再绕环广播一圈，最后每张卡都拿到完整的总和。
为什么要拆成这两个阶段？因为如果走“集中式”的 AllReduce（所有卡把数据发给一张卡，那卡加完再广播回去），总通信量是 O(2N)，且主卡会成为瓶颈。Ring-AllReduce 把通信量压到 O(2N·(P-1)/P)，接近 O(2N)。TP 在工程实现里基本都是用 Ring-AllReduce 或它的变体。
图7里还画了一个进阶版本：把 AllReduce 拆成 ReduceScatter + AllGather 后，如果对数据做进一步“切分”，你并不总是需要 AllGather 那一半——如果下游只需要分片数据，ReduceScatter 的结果就能直接用。这就是“序列并行”的基础思路之一，也是 TP 的通信进一步压榨的切入点。
5.2 一次 Transformer block 的 TP 通信量估算
算一下 TP=2 时一个 Transformer block 的通信量。隐藏维度 H=4096，序列长度 S=2048。
注意力 O 投影的 AllReduce 输入是[S, H]的激活值。每个元素是 2 字节（BF16），所以通信量是 S × H × 2 × 2 = 2048 × 4096 × 4 ≈ 33.6MB。MLP 压缩层的 AllReduce 同样传输[S, H]，又多 33.6MB。一次 block 两次 AllReduce，总通信量约 67MB。
这只是一个 block。一个 32 层的模型，一次前向加反向就是 64 次 AllReduce，总计大约 4.3GB 通信量。如果带宽是 600GB/s（NVLink 4.0 级别），光通信时间就是 7 毫秒；如果带宽是 50GB/s，就是 86 毫秒——这基本上把算力利用率打没了。
所以你会理解为什么 TP 不敢跨节点跑，也理解为什么模型越大通信占比会相对下降（因为计算量是 H^2 级别，而通信量是 H 级别，H 增大时计算/通信比在上升）。
5.3 实测中几个通信优化的手段
省通信的第一招，是“融合”。把 V 个连续的小 AllReduce 合并成一个大 AllReduce。Transformer 里很多位置都能这么干，比如把残差相加和后续 LayerNorm 融合，或者把多个 decoder 层的通信错开。
第二招是“换切分策略”。比如 MLP 的扩展层，如果改成“列切分”，输出虽然要拼一次，但可以把这次拼接跟下一层行切分的 AllReduce 合并成一次组通信。Megatron 的 sequence parallel 就是在这一点上动的手脚，把本来“拼接”的通信变成“沿序列维度切分”，让通信从[S, H]降到[S/P, H]。
第三招是预防性的：避免频繁同步。TP 里 AllReduce 是强同步的，卡与卡之间速度稍有差异，就会互相等待。如果是异构环境，TP 的卡速差异会被放大，尽量保证同一个 TP 组内的卡是同型号、同频率的。
我个人踩过的最深的一个坑就是：把 TP 通信的 AllReduce 默认实现换成了 NCCL 的低延迟模式，结果小模型跑得飞快，一上大模型频繁出现超时。后来排查发现是通信量与 ring 顺序不匹配，改成大数据量模式才稳定。这类参数调优需要实测对比，别照搬别人的配置。
6. 常见问题与排查
6.1 现象速查表
| 现象 | 可能原因 | 排查方向 |
|---|---|---|
| 训练 loss 直接变成 NaN | 权重切分方向错了 | 检查 Column/Row 的切分维度与通信方式是拼接还是相加 |
| loss 正常但吞吐远低于预期 | TP 规模过大或跨节点通信 | 查看 NCCL 日志、通信时间占比，考虑把 TP 换成 PP |
| 显存不够但报错在通信层 | 激活值过大导致峰值显存 | 打开激活检查点，减小 batch size，或把部分层切走 |
| 不同 rank 的 loss 不一致 | AllReduce 没做或做错 | 检查每个 Linear 是否实现了规定的通信语义 |
| 前向正常、反向报梯度形状错误 | 反向传播的梯度切分没处理好 | 核对权重转置后梯度计算对应的维度和 rank 顺序 |
6.2 定位“慢”的实战方法
如果训练跑起来了，但 GPU 利用率很低，先别看是不是 TP 的问题。方法：在每轮迭代里记录每个 Linear 层前后的各 rank 忙碌等待时间，或者直接看nvidia-smi里的吞吐和 SM 占用率。
一个典型场景：TP=8 跑 13B 模型，吞吐只有理论值的 30%。查通信日志发现大量时间花在 AllReduce 的等待上，而算力并没有吃满。这种情况常见原因是 TP 组跨了 PCIe 交换机而不是 NVLink，通信带宽被拓扑限制住了。解法是把 TP 组限制在同一个 NVLink 域内，跨域通信换 PP/DP。
另一个场景：loss 不下降但也没有报错。先跑一个单卡小模型验证数值，再逐步增加 TP 规模，对比输出。大概率是某个模块的切分和通信没有配对——比如前面模块 Row 切分之后，后面模块却想拿完整 X。
6.3 精度问题与累积误差
TP 把一次矩阵乘法拆到多卡再通信，浮点加法的顺序会变，所以结果跟单卡并非完全一致。在大模型长时间训练里，这种误差会累积，尤其是跟 BF16 混合精度叠加时，可能出现 loss 曲线抖动。
我处理过的案例是：同一套配置在 TP=4 时训练稳定，改成 TP=8 后 loss 周期性波动。逐个排查发现是 LayerNorm 的方差统计因为卡间数据分布不一致导致数值抖动。解法是把 LayerNorm 泛化为“残差+LayerNorm+Dropout”一体化融合，并保证跨卡统计顺序一致，抖动消失。
精度问题不一定要完全归零误差，大模型训练本来就允许一定数值差异，关键是差异不能影响收敛方向。如果你发现 TP 规模变大后 loss 曲线明显变差，优先怀疑的是通信或数值稳定性，其次才是学习率策略。
说了这么多，其实最值得动手做的还是两件事：第一，按图3的布局，用第三节的代码在本地跑通一个最小的两层 TP 模型，亲眼看一看每张卡上到底是谁在通信、通信之后长什么样；第二，找一个你熟悉的小模型，试着把它的一个 Linear 层改成 TP 实现，然后对比单卡输出误差和显存变化。我自己是这么学明白的，很多东西看十遍图不如动手改一遍代码来得透彻。TP 并不玄乎，本质就是“把一个大矩阵乘法拆成几个小矩阵乘法，再把该拼的拼起来、该加的加起来”，多看几张图、多写几行代码，它就会变成你工具箱里很顺手的一件工具。
