如果你动手训练过一个稍大规模的文本生成模型，大概率遇到过这样的场景：模型参数做到几十亿，词表开到五万，为了照顾小说章节的上下文把序列长度拉到2048，然后单卡训练一启动就OOM。就算勉强塞进去，一次迭代也要跑几十秒，照这个速度去预训练一个像样的续写模型，时间成本根本扛不住。
要解决小说大模型的分布式训练问题，张量并行是绕不开的一环。它是目前主流大模型训练里最基础的“分而治之”手段之一，把单个矩阵运算切成多份交给多张卡协作完成，直接降低单卡显存压力、提升单步吞吐。这篇文章我会从原理、通信开销、代码实现到踩坑排错，把张量并行这条线完整讲透。适合正在做大模型训练、想自己搭分布式环境，或者被训练资源卡住的工程师和研究者参考。
1. 小说大模型为什么必须分布式训练：显存账本与并行方案选型
1.1 显存账本：不只是权重，还有一堆“附加品”
很多人对训练显存的第一反应是“模型多大就占多大”，实际上训练态的显存开销远不止权重。以常见做法为例，参数用FP16或BF16存储，每个参数占2字节，这部分叫模型权重。但训练时我们还必须存梯度，每个参数同样至少2字节。更夸张的是优化器状态，很多项目用AdamW，它会为每个参数保存三份FP32数据：主权重、一阶动量、二阶动量，也就是每个参数12字节。算下来，训练一个35亿参数的中等模型，理论显存需求就是35B×(2+2+12)字节，约560GB。单卡主流显卡也就80GB上下，HPC级别的卡也一样装不下。
小说模型还有一个特别吃亏的地方：词表大。为了保证生成词汇丰富，词表动辄五万、十万，Embedding层的参数量会非常夸张。词表五万、隐藏维度4096，光输入Embedding就是2亿参数，再加上输出层一般还权重共享，这部分直接摊进每次前反向里。所以真实项目里，同样的参数量，文本模型往往比视觉模型更吃显存。
另一个隐性因素是激活值。序列长度2048、批次4、隐藏维度4096时，每一层Transformer的中间激活都是GB量级，几十层叠下来非常可观。这也是为什么哪怕你把权重硬塞进显存，训练起来还是容易OOM。
1.2 三大并行范式的分工
分布式训练里，我习惯先把问题拆成两种：一种是一张卡已经装不下完整模型，另一种是单卡能装但算得太慢、吞吐不够。按这个思路，主流的并行范式可以这样归类：
| 并行方式 | 核心思路 | 解决的主要问题 | 通信代价 |
|---|---|---|---|
| 数据并行 | 每张卡一份完整模型副本，切片喂不同batch | 吞吐不足、算力利用不够 | 每步同步梯度，通信量大 |
| 张量并行 | 把单个层内的大矩阵按行/列切分，多卡协作算一个GEMM | 单卡放不下权重、单层计算成为瓶颈 | 每层前反向多次AllReduce，通信极频繁 |
| 流水线并行 | 按层切分，每张卡负责若干层 | 激活值峰值过高、超深模型装不下 | 仅在切分边界传递激活与梯度，通信量相对小 |
实际项目几乎不会只用一种。最常见的组合是“数据并行 + 张量并行 + 流水线并行”三种混用，数据并行照顾整体吞吐，张量并行解决单层矩阵太大，流水线并行解决模型层数太深。
1.3 张量并行适合什么场景，又有什么边界
我的判断标准很简单：当单层矩阵的参数量和计算量已经成为瓶颈时，优先考虑张量并行。Transformer里最典型的就是Attention的QKV投影和MLP的两层线性层，这几处每个都是超大矩阵，张量并行把它们切碎后，单卡显存和单卡计算压力都能明显下降。
但张量并行的前提是通信非常快。它每一步迭代都要在TP组内做AllReduce，通信量在GB级别。如果两张卡之间只有普通千兆网卡，那训练速度会被通信死死按在地上。所以实际工程里，张量并行几乎只在单机内使用，依赖NVLink或PCIe Switch。跨节点做张量并行，除非你有极高带宽的专业网络，否则性能一定很难看。这也是为什么很多框架默认把TP大小限制在8以内，超过8就要认真评估收益。
2. 张量并行核心原理：矩阵切分与通信算子的设计思路
2.1 矩阵切分的两种基本姿势
要理解张量并行，先看一个最简单的全连接层 Y = X·W。假设输入X形状是[batch, in_features]，权重W形状是[in_features, out_features]，我们有两种切法。
第一种叫列并行。把W按输出维度切成左右两块，每张卡只保留一块，输入X保持完整副本，那么每张卡算出来的是部分输出，最后把多张卡的输出沿特征维拼接起来，就得到完整结果。这种切法很适合“输出变大”的场景，比如词表投影。
第二种叫行并行。把W按输入维度切成上下两块，同时把输入X也沿特征维切碎，每张卡只拿到输入的一部分和权重的一部分，各算各的部分和，最后用一次AllReduce把多张卡的部分和相加，得到最终输出。
生活化类比：好比要算一个超长算式，列并行是每个人算不同几列，最后拼表；行并行是每个人算不同几行的部分结果，最后汇总求和。
2.2 MLP层的切分策略：f算子与f-bar算子的妙处
Transformer里MLP通常是两个线性层夹一个激活函数：先升维，再过GeLU，再降维。对这两个线性层怎么切，直接决定通信次数。
标准做法是第一个线性层用列并行，第二个线性层用行并行。为什么这样搭？因为GeLU是对每个元素独立操作的激活函数，它“不认识”其他位置的数据。第一个线性层输出被切分后，每张卡在自己的分片上做GeLU，完全不需要通信。然后第二个线性层是行并行，它自然会把多张卡的部分结果相加，这里做一次AllReduce，整层的输出就完整了。
这里有两个特殊算子，我简单整理一下：
| 算子 | 前向是否通信 | 反向是否通信 | 典型用途 |
|---|---|---|---|
| f算子 | 不通信，本地计算列并行输出 | 需要AllReduce，累加所有卡对输入的梯度 | 列并行线性层 |
| f-bar算子 | 需要AllReduce，把各卡部分和相加 | 不通信，各卡本地计算反传梯度 | 行并行线性层 |
这两个算子最妙的地方，是把通信均匀摊到了前向和反向。f层前向不通信，省掉一次等待，代价是反向必须AllReduce；f-bar层正相反。整体看，每一层MLP在前反向各一次AllReduce，通信次数不多不少，设计非常紧凑。
如果两个线性层都切成列并行，中间就需要一次拼接通信，输出端还得再来一次拼接，等于平白多了一次传输。如果都切成行并行，输入切分和输出切分的配合又很难受。所以工程上很默契地固定成“列并行 + 行并行”。
2.3 自注意力层的切分：多头天然适合张量并行
自注意力部分比MLP更特殊，因为它本身就是按“头”设计的。Q、K、V三个投影矩阵，完全可以按输出维度切分，让每张卡只计算一部分注意力头。具体来说，假设TP大小为N，原本有H个注意力头，那么每张卡负责H/N个头。卡在自己的分片上做QKV投影，然后本地算注意力得分、加权求和，得到部分头的输出。最后接一个行并行的输出投影层，用f-bar做一次AllReduce把各卡结果合并。
这种切法几乎是为多头注意力量身定做的，通信点也清晰：QKV投影之间无通信，注意力计算无通信，只在输出投影结束做一次AllReduce。我们不需要额外设计“跨卡注意力”，每张卡都是独立算自己的头，只是头变少了。
需要留意，Attention输入X在每张卡上都必须是完整副本。因为每张卡都要用自己的Q去和完整的K、V做注意力。如果上一层输出是切分的，必须先AllReduce恢复完整，再进入注意力层。这也是张量并行实现里最容易混淆的地方。
2.4 不切的层：LayerNorm和Embedding怎么处理
LayerNorm和Dropout在张量并行里通常直接复制到每张卡，不做切分。原因是它们要么是逐元素操作，要么参数量小到可以忽略，切了反而要引入额外的通信同步点。LayerNorm有两个参数gamma和beta，形状都是[hidden_size]，对动辄上亿的参数来说只是零头，复制一份更省事。
Embedding则比较纠结。小说模型词表五万到十万，Embedding参数量可能占整个模型的20%以上，看着很想切。但切Embedding有个麻烦：它既是输入层又是输出层，输出概率需要收集所有卡的logits做AllGather，反向时又要再分散梯度，通信开销不小。我的建议是项目初期先复刻Embedding，保持正确性和简单性，等整个训练流程稳定了，再考虑词表并行的优化方案。有时候模型总参数不大，复刻Embedding占的那点显存，完全可以用其他方式补回来。
3. 通信开销分析与架构调优：别让带宽拖垮训练
3.1 一次训练迭代的通信量到底有多大
很多人问：张量并行不是把计算分开了吗，为什么训练速度没有线性提升？答案就在通信里。每层MLP要做两次AllReduce，每层Attention输出还要一次，几十层叠加，通信总量非常惊人。
我做过一个粗略估算：假设模型隐藏维度4096，层数32，序列长度2048，批次大小8，TP大小8。只看MLP部分，每一步迭代的通信量大约等于 层数 × 批次 × 序列长度 × 隐藏维度 × 2（前向一次加反向一次）再乘2字节的BF16精度。代入数字就会发现，单卡参与同步的数据量已经到GB级别，这还只是MLP，没算Attention和Embedding。
| 阶段 | 通信内容 | 每次通信数据量 |
|---|---|---|
| MLP前向 | f-bar层的AllReduce输出 | batch × seq × hidden |
| MLP反向 | f层的AllReduce输入梯度 | batch × seq × hidden |
| Attention输出投影 | 行并行的AllReduce | batch × seq × hidden |
| 其他 | 必要时广播输入、同步梯度 | 视实现而定 |
所以张量并行对硬件的要求极高。TP组内必须是高速拓扑，最理想的是单机内8卡NVLink全互联。即使是PCIe Switch，通信带宽也远低于NVLink，训练效率会明显打折。这也解释了为什么张量并行不适合盲目扩大规模。
3.2 标准AllReduce与序列并行的取舍
早期张量并行实现里，所有通信都走AllReduce，逻辑简单，正确性容易保证。但有一批研究者发现，LayerNorm和Dropout前的输入其实可以不保留完整副本，把输入按序列维度切分，只在需要全量特征时再做通信，这就是序列并行。序列并行能把每层输入X的激活显存从全局完整副本降为序列分片，同时把MLP的两次通信降为一次。
实际训练小说模型时，序列长度通常很长，2048都算短的，甚至有人上4096或8192。序列并行对显存的改善非常明显。很多成熟框架里已经把张量并行和序列并行打包在一起，你开了TP开关，序列并行往往自动跟着开。我的建议是：只要框架支持，序列并行能开就开，性价比非常高。
3.3 通信计算重叠：让卡别“干等”数据
分布式训练最怕的就是计算和通信串行：先算一阵，然后停下来等AllReduce，再接着算。解决思路是把通信藏进计算时间里。最常用的手段是micro-batch切分。假设全局批次大小是32，我们不一次算完，而是切成4个micro-batch，每个8条样本。第1个micro-batch算完发起通信，在等待结果的同时，GPU立刻开始计算第2个micro-batch，通信和计算就叠起来了。这就是流水线式重叠。
现代训练框架里，这个能力往往通过异步通信接口和CUDA Stream实现。我个人做性能排查时，会先用NCCL的timeline或者PyTorch Profiler看每个算子的耗时，重点看AllReduce前后有没有明显“气泡”。如果通信占比超过30%，就要考虑增大micro-batch数量、减少同步频率，或者换更高带宽的互联方案。
3.4 什么时候开TP，开到多大
有一个很常见的误区：一上来就追求张量并行规模。其实TP不是越大越好。当模型权重能完整放进单卡显存时，开TP只会增加通信开销，完全没有必要；单卡跑得动，优先用纯数据并行。只有当单卡装不下，或者单卡训练慢到无法接受时，才需要引入TP。
我一般在项目里的判断顺序是这样：
- 先算显存账：权重、梯度、优化器状态、激活值分别预估。
- 单卡能装且吞吐可接受，用数据并行加ZeRO优化器状态。
- 单卡装不下，优先TP=2或TP=4，配合序列并行和激活重计算。
- TP超过8之前，先认真评估硬件拓扑。跨节点的TP，除非有InfiniBand级别的互联，否则果断排除。
TP大小还要和模型结构对齐：TP=4时，Attention的头数必须能被4整除，hidden_size也要能整除。这个在设计模型时就得考虑，不然后面改起来非常痛苦。
4. 从零实现张量并行训练：代码骨架与关键步骤
4.1 环境准备与TP组初始化
如果用PyTorch生态，最常用的启动方式是torchrun。假设单机8卡，TP_SIZE为4，数据并行大小就是2。我习惯这样组织：同一节点内rank连续，先按TP_SIZE切组，组内用NVLink高速通信。
torchrun --nproc_per_node=8 --nnodes=1 train.py --tp_size 4初始化进程组后，要手动建一个TP子组。这里的核心逻辑是：每个rank先算自己属于哪一个TP组，再把组内rank列表传给new_group。
import torch.distributed as dist dist.init_process_group(backend="nccl") world_size = dist.get_world_size() rank = dist.get_rank() tp_size = 4 tp_rank = rank % tp_size tp_group = dist.new_group(ranks=[r for r in range(world_size) if r % tp_size == tp_rank])这个写法假设rank号就是按节点连续排的。实际多机时，不同节点的rank会交错，建议把节点内rank映射到连续索引，再切TP组，避免跨节点组内通信。
4.2 核心模块：ColumnParallelLinear与RowParallelLinear
理解了原理，实现其实不难。关键在于自定义的线性层要“知道自己该通信哪一次”。我给出一个可运行逻辑的简化版本，实际工程会加缓存、半精度处理等优化，但骨架就是这样。
import torch import torch.nn as nn import torch.nn.functional as F import torch.distributed as dist class ColumnParallelLinear(nn.Module): """列并行线性层：切输出维度，前向无通信，反向做AllReduce""" def __init__(self, in_features, out_features, tp_size, tp_rank): super().__init__() self.tp_size = tp_size self.tp_rank = tp_rank local_out = out_features // tp_size self.weight = nn.Parameter(torch.empty(local_out, in_features)) nn.init.kaiming_uniform_(self.weight, a=5 ** 0.5) def forward(self, x): # x是完整输入，每张卡各自算本地部分 local_out = F.linear(x, self.weight) # 反向时需要AllReduce对x的梯度，这里用钩子实现 x.register_hook(lambda grad: reduce_grad(grad)) return local_out class RowParallelLinear(nn.Module): """行并行线性层：切输入维度，前向做AllReduce，反向无通信""" def __init__(self, in_features, out_features, tp_size, tp_rank): super().__init__() self.tp_size = tp_size self.tp_rank = tp_rank local_in = in_features // tp_size self.weight = nn.Parameter(torch.empty(out_features, local_in)) nn.init.kaiming_uniform_(self.weight, a=5 ** 0.5) def forward(self, x): local_out = F.linear(x, self.weight) dist.all_reduce(local_out, group=TP_GROUP) return local_out这里需要补充一个reduce_grad函数：对列并行层来说，前向不做通信，但反向时每张卡都会算出一份对输入X的局部梯度，必须AllReduce把所有局部梯度和起来，才能得到完整梯度，否则模型更新会错乱。
def reduce_grad(grad): dist.all_reduce(grad, group=TP_GROUP) return grad提示：注册钩子实现AllReduce在功能上没问题，但工程上效率不高。成熟的框架会把“前向不通信、反向通信”封装成自定义autograd.Function，避免钩子带来的调度开销。追求性能时建议照这个方向改造。
4.3 把Transformer层拆成并行“积木”
有了两个基础模块，改造Transformer层就变成“替换零件”的操作。一个典型的Transformer层包含Attention和MLP两部分，我给出替换后的伪代码结构。
class ParallelAttention(nn.Module): def __init__(self, hidden, heads, tp_size, tp_rank): super().__init__() self.qkv = ColumnParallelLinear(hidden, hidden * 3, tp_size, tp_rank) self.out_proj = RowParallelLinear(hidden, hidden, tp_size, tp_rank) def forward(self, x, mask=None): # 列并行QKV，本地算部分头 qkv = self.qkv(x) # 切分并 reshape 成多头 # 本地计算scaled dot-product attention # 得到部分头的输出 attn_out = row_parallel_reshape(...) return self.out_proj(attn_out)MLP部分则是两个模块串联：
class ParallelMLP(nn.Module): def __init__(self, hidden, ffn_hidden, tp_size, tp_rank): super().__init__() self.fc1 = ColumnParallelLinear(hidden, ffn_hidden, tp_size, tp_rank) self.fc2 = RowParallelLinear(ffn_hidden, hidden, tp_size, tp_rank) def forward(self, x): x = self.fc1(x) x = F.gelu(x) return self.fc2(x)每一层Transformer输出前，记得把LayerNorm放在最前面做PreNorm，输入X保持完整副本。TP改造后，模型内部通信已经由线性层的AllReduce承担，外部看不到额外操作。
4.4 训练循环、超参数与性能验证
训练循环本身并不需要太多改动，依然按常规方式算loss、反传、更新优化器。只是有几个点必须注意。
一是loss只在一个rank上算即可，不需要所有rank都同步一份logits再算。如果你不做词表并行，输出层是完整复制的话，loss计算跟单卡完全一样。二是优化器更新时，每个rank只更新自己那份分片权重，不做额外通信，更新完各卡的分片自然组成完整模型。三是最开始调试时，强烈建议先用一个小模型做“正确性对照”：在单卡上跑同样的结构，在TP=2或TP=4上跑，对比前几步loss和梯度。数值不会完全一样，因为浮点累加顺序变了，但loss曲线趋势应该一致。
超参数方面，我的经验是TP本身不改变每步的有效批次大小，但它会改变单卡显存占用上限，所以你可以尝试调大micro-batch。在小说模型场景下，我通常把序列长度设2048，micro-batch从1开始小步往上探，遇到OOM再往下调。学习率还是沿用warmup加余弦衰减，常用初始区间是1e-4到3e-4，具体看模型规模。另外强烈建议开启激活重计算，也就是gradient checkpointing，它会把每层激活值丢弃，反传时重新计算，虽然多花一点算力，但显存能省下很大一块。
性能上，我最看重的指标是tokens/s，也就是每秒处理的token数。记录一次完整迭代耗时，乘以批次大小和序列长度，再除以卡数，就知道单卡吞吐。有了这个指标，调参才有依据。不要只看显存占用，显存省下来了但速度慢得像蜗牛，一样没有意义。
5. 实战排错手册：显存、通信、收敛的九九八十一难
5.1 显存不降反增的迷思
一个很反直觉的现象：明明把权重切开了，显存却比单卡还紧张。这种情况多半是激活值在作怪。张量并行切的是权重和梯度，但很多层（比如Attention的输入X）仍然是完整副本，每张卡都存一份，激活显存不仅没降，还可能因为中间结果分片存储导致额外开销。
我的排查路径很清楚：先用torch.cuda.max_memory_allocated()打印峰值显存，再用显存分析工具看每一层占用，区分是激活还是权重。如果是激活太大，三条路：开激活重计算，开序列并行，或者减小micro-batch。小说模型序列很长，激活重计算基本是必需品。
5.2 通信卡死或慢到离谱
如果训练一开始就卡住不动，十有八九是通信问题。先看日志，把NCCL_DEBUG=INFO打开，NCCL会打印它检测到的拓扑和使用的通信协议。再用nvidia-smi topo -m看卡间连接，确认TP组都在同一个NVLink域里。如果发现TP组里混了跨PCIe Switch甚至跨CPU的卡，通信性能会一落千丈。
还有一个容易被忽略的坑：NCCL的P2P传输有时会因为驱动或固件问题失败，此时通信会回退到共享内存或网络，速度同样很难看。可以试一下NCCL_P2P_DISABLE=1环境变量，强制走其他路径，对比耗时。如果开了之后反而变快，基本就是P2P没走对。
注意：TP通信是高频小消息，单次AllReduce的数据量并不大，但对延迟极敏感。跨节点的网络延迟通常是NVLink的几十倍，所以除非万不得已，不要跨节点组TP。
5.3 梯度对不上、loss不降怎么定位
训练能跑但loss不降，大概率不是训练策略问题，而是并行实现有bug。我见过最多的是两种情况：第一，列并行层反向忘了AllReduce输入梯度，导致每张卡更新时看到的梯度不完整；第二，AllReduce用错了group，把TP组和全局组混在一起，导致梯度被重复累加很多次。
排查方法其实不难。先固定随机种子，在TP=1和TP=4下各跑几步，打印每层梯度范数。如果TP=4的梯度范数和TP=1差太多，基本就是通信算子放错了位置。再用torch.autograd.detect_anomaly打开异常检测，虽然慢，但能更快定位到NaN或梯度爆炸的源头。
5.4 模型保存与恢复：分片权重怎么处理
张量并行下，模型权重天然就是分片状态：每张卡只保存自己那份，直接torch.save(model.state_dict())会把分片结构也存下来，等到改变TP大小恢复时就会报shape不匹配。我的经验是保存时按rank分文件，文件名里带上tp_rank和dp_rank，同时额外保存一份模型的配置信息，包括tp_size、层数、隐藏维度、词表大小。
切换TP规模时，需要写一个转换脚本，把旧的分片权重合并成完整权重，再按新的TP配置重新切分。这个脚本要提前写好，不然等训练到一半想调整架构，会发现卡在模型加载上。优化器状态也是同样思路，保存时带上流水线和张量并行的切分元信息。
5.5 小说模型训练特有的细节
最后聊几个跟文本生成强相关的坑。第一，采样验证阶段如果用了词表并行，必须记得AllGather所有卡的logits，直接拿分片的logits做argmax会得到错误token。第二，小说类数据经常出现超长段落，训练时要做好长度截断和padding策略，我一般用动态padding，同一个batch内尽量对齐，减少无效token的浪费。第三，中文小说分词差异很大，词表选太小会出现大量未登录词，选太大会吃掉太多显存，这个取舍要前置到模型设计阶段。
6. 最后几句实在话
从我的实际经验看，张量并行真正难的从来不是“把矩阵切开”，而是把通信、显存、计算三者调到一个平衡点。我见过太多人一上来就开TP=8，结果通信开销比计算还大，训练速度反而不如TP=4配合激活重计算。硬件拓扑、模型结构、序列长度、batch大小，每个因素都会改变最优配置，别人能用得好，不代表你的环境也能照搬。
所以我特别建议：项目启动时从TP=1的单卡基准确认正确性，再逐步扩展到TP=2、TP=4，每一步都记录吞吐和显存。小步走，每一步都有对照，出了问题也容易定位。真正踩过几次坑之后你会发现，分布式训练里最值钱的不是某一个花哨算子，而是一套能在不同规模下稳定复现、可排查、可迁移的训练路线。希望这篇内容能帮你在自己的小说大模型项目里少走几步弯路。
