张量并行实战：小说大模型分布式训练的原理、通信与调优

如果你动手训练过一个稍大规模的文本生成模型大概率遇到过这样的场景模型参数做到几十亿词表开到五万为了照顾小说章节的上下文把序列长度拉到2048然后单卡训练一启动就OOM。就算勉强塞进去一次迭代也要跑几十秒照这个速度去预训练一个像样的续写模型时间成本根本扛不住。要解决小说大模型的分布式训练问题张量并行是绕不开的一环。它是目前主流大模型训练里最基础的“分而治之”手段之一把单个矩阵运算切成多份交给多张卡协作完成直接降低单卡显存压力、提升单步吞吐。这篇文章我会从原理、通信开销、代码实现到踩坑排错把张量并行这条线完整讲透。适合正在做大模型训练、想自己搭分布式环境或者被训练资源卡住的工程师和研究者参考。1. 小说大模型为什么必须分布式训练显存账本与并行方案选型1.1 显存账本不只是权重还有一堆“附加品”很多人对训练显存的第一反应是“模型多大就占多大”实际上训练态的显存开销远不止权重。以常见做法为例参数用FP16或BF16存储每个参数占2字节这部分叫模型权重。但训练时我们还必须存梯度每个参数同样至少2字节。更夸张的是优化器状态很多项目用AdamW它会为每个参数保存三份FP32数据主权重、一阶动量、二阶动量也就是每个参数12字节。算下来训练一个35亿参数的中等模型理论显存需求就是35B×(2212)字节约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层前向不通信省掉一次等待代价是反向必须AllReducef-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批次大小8TP大小8。只看MLP部分每一步迭代的通信量大约等于 层数 × 批次 × 序列长度 × 隐藏维度 × 2前向一次加反向一次再乘2字节的BF16精度。代入数字就会发现单卡参与同步的数据量已经到GB级别这还只是MLP没算Attention和Embedding。阶段通信内容每次通信数据量MLP前向f-bar层的AllReduce输出batch × seq × hiddenMLP反向f层的AllReduce输入梯度batch × seq × hiddenAttention输出投影行并行的AllReducebatch × 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优化器状态。单卡装不下优先TP2或TP4配合序列并行和激活重计算。TP超过8之前先认真评估硬件拓扑。跨节点的TP除非有InfiniBand级别的互联否则果断排除。TP大小还要和模型结构对齐TP4时Attention的头数必须能被4整除hidden_size也要能整除。这个在设计模型时就得考虑不然后面改起来非常痛苦。4. 从零实现张量并行训练代码骨架与关键步骤4.1 环境准备与TP组初始化如果用PyTorch生态最常用的启动方式是torchrun。假设单机8卡TP_SIZE为4数据并行大小就是2。我习惯这样组织同一节点内rank连续先按TP_SIZE切组组内用NVLink高速通信。torchrun --nproc_per_node8 --nnodes1 train.py --tp_size 4初始化进程组后要手动建一个TP子组。这里的核心逻辑是每个rank先算自己属于哪一个TP组再把组内rank列表传给new_group。import torch.distributed as dist dist.init_process_group(backendnccl) 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, a5 ** 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, a5 ** 0.5) def forward(self, x): local_out F.linear(x, self.weight) dist.all_reduce(local_out, groupTP_GROUP) return local_out这里需要补充一个reduce_grad函数对列并行层来说前向不做通信但反向时每张卡都会算出一份对输入X的局部梯度必须AllReduce把所有局部梯度和起来才能得到完整梯度否则模型更新会错乱。def reduce_grad(grad): dist.all_reduce(grad, groupTP_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, maskNone): # 列并行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只更新自己那份分片权重不做额外通信更新完各卡的分片自然组成完整模型。三是最开始调试时强烈建议先用一个小模型做“正确性对照”在单卡上跑同样的结构在TP2或TP4上跑对比前几步loss和梯度。数值不会完全一样因为浮点累加顺序变了但loss曲线趋势应该一致。超参数方面我的经验是TP本身不改变每步的有效批次大小但它会改变单卡显存占用上限所以你可以尝试调大micro-batch。在小说模型场景下我通常把序列长度设2048micro-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_DEBUGINFO打开NCCL会打印它检测到的拓扑和使用的通信协议。再用nvidia-smi topo -m看卡间连接确认TP组都在同一个NVLink域里。如果发现TP组里混了跨PCIe Switch甚至跨CPU的卡通信性能会一落千丈。还有一个容易被忽略的坑NCCL的P2P传输有时会因为驱动或固件问题失败此时通信会回退到共享内存或网络速度同样很难看。可以试一下NCCL_P2P_DISABLE1环境变量强制走其他路径对比耗时。如果开了之后反而变快基本就是P2P没走对。注意TP通信是高频小消息单次AllReduce的数据量并不大但对延迟极敏感。跨节点的网络延迟通常是NVLink的几十倍所以除非万不得已不要跨节点组TP。5.3 梯度对不上、loss不降怎么定位训练能跑但loss不降大概率不是训练策略问题而是并行实现有bug。我见过最多的是两种情况第一列并行层反向忘了AllReduce输入梯度导致每张卡更新时看到的梯度不完整第二AllReduce用错了group把TP组和全局组混在一起导致梯度被重复累加很多次。排查方法其实不难。先固定随机种子在TP1和TP4下各跑几步打印每层梯度范数。如果TP4的梯度范数和TP1差太多基本就是通信算子放错了位置。再用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. 最后几句实在话从我的实际经验看张量并行真正难的从来不是“把矩阵切开”而是把通信、显存、计算三者调到一个平衡点。我见过太多人一上来就开TP8结果通信开销比计算还大训练速度反而不如TP4配合激活重计算。硬件拓扑、模型结构、序列长度、batch大小每个因素都会改变最优配置别人能用得好不代表你的环境也能照搬。所以我特别建议项目启动时从TP1的单卡基准确认正确性再逐步扩展到TP2、TP4每一步都记录吞吐和显存。小步走每一步都有对照出了问题也容易定位。真正踩过几次坑之后你会发现分布式训练里最值钱的不是某一个花哨算子而是一套能在不同规模下稳定复现、可排查、可迁移的训练路线。希望这篇内容能帮你在自己的小说大模型项目里少走几步弯路。