企业官网建设流程全解析

做训练优化的第一课，往往不是调参，而是面对“显存不够”四个字时选哪条路。我第一次搭多卡训练脚本，先是让模型可以在单卡上摊开，然后发现要摊开的不仅是权重，还有梯度和优化器状态，最后对着一块几十GB的激活张量发呆。那段时间最常被我挂嘴边的方案就是TP，也就是张量并行（Tensor Parallelism）。当时我对它的理解止步于“把矩阵切开，分给几张卡”，等到真的把切分规则落到Transformer上时，才意识到“怎么切、切完怎么拼、拼完怎么通信”才是整个方案的主体。这篇文章，我尽量用切分示意图的方式，把TP的完整链路从矩阵到底层通信、再到和各类并行策略的组合讲清楚，最后附上自己工程实践里踩过的真实坑。

1. 一张卡放不下模型，为什么先想到的是张量并行

1.1 显存花在哪，四个口袋说完

先纠正一个常见的误解：训练大模型时，显存不全是模型权重占掉的。一个超过十亿参数的中型模型，用半精度存放权重，本身就要两个多GB；到十几B规模，二十多GB起步。这个数字很多人都有基本概念，但别忘了梯度还要一份同样大小的空间。更夸张的是优化器状态——如果用了自适应矩估计这类主流优化器，每个参数要维护两份额外的全精度状态，加起来经常是参数总量的数倍。

也就是说，光权重、梯度、优化器状态这三件套，就足以把一块显存不算小的加速卡直接压穿。这还没算激活值。前向计算时，每个线性层、每次注意力计算都会把中间结果留在显存里供反向传播使用。序列越长、batch越大、层越深，这部分就越膨胀。训练长序列时，激活值经常比模型参数本身还大。这也是为什么同一个模型做推理可能很轻松，训练却怎么都塞不进单卡的原因。

面对这种压力，数据并行救不了。数据并行让每张卡各处理自己的数据，再把梯度做归约求和，听着很合理，可前提是每张卡都要有一份完整模型副本。模型放不下，数据并行就只能干看着。而张量并行（TP）的做法完全不同：它的切割对象是单个权重矩阵本身，把层内部切分，而不是把整层搬走。所以TP是“层”这个粒度以下最直观、最细粒度的武器。

1.2 为什么不先考虑流水线并行

有人会说，显存不够可以按层切啊，这就是流水线并行。流水线把网络按层切成若干段，每张卡只负责其中一段网络。它确实能显著降低每张卡存的权重数量，在跨节点分布式训练里也很常用。但流水线并行有两个先天问题：第一，模型足够大时，单个Transformer层内部的参数和激活也可能大到放不下，流水线帮不了你；第二，流水线天然存在气泡，计算单元在流水槽位空转的时间要靠精细的微批次调度去填补，效率天花板肉眼可见。

更关键的是，流水线并行对激活值的切分并不友好。层与层之间传递的激活张量，在到达下一张卡时基本是整块移动的，单层内部的巨大中间结果依然无法拆开。而TP是在单个矩阵乘法的维度上动手，权重、梯度、优化器三件套全都能跟着切分，激活值的一部分也能顺势拆开，是破解超单卡容量最直接的方式。

在我搭分布式训练脚本的初期，最让我头疼的恰恰是激活值。因为参数、梯度和优化器都有比较固定的存储占用，算一算就能规划清楚，而激活值随序列长度膨胀的速度几乎失控。TP恰好能同时管住模型的存储和一部分中间结果，这就是我每次都先考虑它的原因。

2. 先看懂矩阵切分：张量并行的两条基本路径

2.1 列并行：每张卡拿到一窄条权重，输出靠拼接拼回来

一个线性层的计算是Y = XW。X的形状是（batch, seq, H_in），W的形状是（H_in, H_out）。列并行选择的切割位置是W的输出维，也就是H_out这一侧。把W看成一张大表，列并行就像拿一把竖切刀，把输出通道切成N段窄条，每张卡分到一窄条权重。示意图如下：

W (H_in × H_out) → 竖切N块 卡0: W[:, 0 : H_out/N] → Y0 = X · W0 卡1: W[:, H_out/N : 2H_out/N] → Y1 = X · W1 ... 卡N-1: W[:, (N-1)H_out/N : H_out] → Y(N-1) = X · W(N-1) 最终 Y 沿输出维拼接：Y = [Y0, Y1, ..., Y(N-1)]

这里输入X是完整复制的，相当于每张卡都拿到了原来那个输入矩阵。由于每张卡只保留了一窄条权重，矩阵乘法里的计算量降为原来的1/N，对应的权重和梯度存储也各自只占1/N。但代价是，要还原成完整输出Y时，各卡的结果只是H_out维度上的一个片段：卡0有前H_out/N列，卡1有后H_out/N列，必须用一次AllGather操作把N份片段沿着输出维拼起来。

反向传播时，完整Y的梯度也需要按同样的位置切回去，让每张卡只拿到自己的那一窄条。整体趋势很清晰：列并行中，输入需要全量，输出需要拼接。这个切法的好处是直观好写，坏处是如果上一层本来就是行并行，AllGather和AllReduce交替起来，通信环节的叠加会让显存和带宽一起打架。

2.2 行并行：输入先分家，输出靠求和凑齐

行并行换一个角度切W，这次从输入维H_in下手，把W横着切成N块，每一块的行数是H_in/N。同时，输入X的最后一维也要按同样比例切分。X被拆成N份，每张卡拿着自己对应位置的X分片和W行块，算一个部分和。示意图如下：

X (batch, seq, H_in) → 按 H_in 切N份 卡0: X[:, :, 0 : H_in/N] · W[0 : H_in/N, :] → Y0 卡1: X[:, :, H_in/N : 2H_in/N] · W[H_in/N : 2H_in/N, :] → Y1 ... 卡N-1: X[:, :, (N-1)H_in/N : H_in] · W[(N-1)H_in/N : H_in, :] → Y(N-1) 最终 Y = Y0 + Y1 + ... + Y(N-1) ← 逐元素求和

矩阵乘法本身有“按行分块后相加”的性质。每张卡持有的部分和形状恰好是完整的（batch, seq, H_out），只是数值上都缺了其他卡那一部分。把N个部分和逐元素相加，就能得到正确结果，标准的AllReduce归约求和路由刚好干这件事。

所以行并行和列并行的通信方向刚好相反：行并行输入要切分、输出要归约求和；列并行输入要完整、输出要拼接。有趣的是它们俩在张量形状上完全互补——行并行吃不消“输入是完整复制”的场景，但一旦上一层是列并行，上一层拼接前的输出切片天然就是这一层的输入分片，能省掉一次通信。

2.3 两条路径的真正关系：它们是一对互补切分

把两个并行放在一起看，第一直觉容易把它们理解成两个独立技巧，其实它们是搭档。Transformer的层结构里，每个子层基本上由“投影进来、算一算、投影出去”构成。一个自然模式是：第一层投影用列并行，把输入完整分发、输出切开；中间的逐点运算按各自的张量位置处理；第二层投影用行并行，接收上一层的输出切片当输入，最后用AllReduce把结果合拢。这样一来，层与层之间的分布状态是连续的，不需要反复把整个张量变回完整状态。

举个例子，前置的列并行输出是沿H_out维拼接的N份。行并行正好要求输入X按H_in维切成N份。如果H_in恰好等于上层的H_out，那么这两种切分天然对齐，行并行可以直接消费列并行吐出的切片，不用额外AllGather去凑完整输入。这就构成了后面Transformer部分反复引用的配对结构。

写到这里，顺便给一个经验结论：不要单独去记哪个并行“更好”，要问自己当前这一步的输入是什么分布模式、下一步的输出需要什么分布模式。TP的设计过程本质上是在做分布规划，所有算子和切分纬度的选择，都是在“完整”与“分片”两种状态之间做路由。

3. Transformer满血版：把TP落到注意力与MLP上

3.1 自注意力模块：QKV、多头、输出投影的切法

注意力子层由三部分组成：QKV投影、注意力计算、输出投影。TP的实际排布，就是围绕“多头”和“拼接”做文章。

先看QKV投影。输入X通常已经是一种切分状态。在经典TP模式里，QKV投影采用列并行：把Q、K、V三个投影矩阵的输出维按比例切到N张卡上。这里的切分维度是3H，也就是QKV合并后的输出通道数。但实践中更聪明的做法是沿着注意力头数去切。假设模型共head_num个头，TP规模为N，就要求head_num能被N整除，然后每卡分到head_num/N个完整注意力头。每个卡上的QKV投影只计算自己那部分头的输出，注意力计算完全隔离在各卡的本地头上，不需要跨卡同步。

这个过程可以画成下面这样：

X → 卡0: QKV投影(只属于自己那部分头) → Q0, K0, V0 → 卡1: QKV投影(头不同) → Q1, K1, V1 ... 每个卡在本地上完成多头注意力中的部分头 → 卡i 输出的是自己那 head_num/N 个头的上下文向量

由于自注意力必须看到一个序列的所有token，序列维度在每个卡上都是完整存在的。这样做，注意力矩阵的大小不受TP影响，只受每卡负责的头数影响。TP=2与TP=8在注意力阶段，单卡的计算量差异就是头数减少带来的。

接下来是输出投影。输出投影的输入是所有注意力头输出沿hidden维拼接好的向量。在TP状态下，各卡已经天然拥有自己那head_num/N个头的输出，把它们在本地拼起来，就得到一个hidden维的局部向量。恰好输出投影可以用行并行的方式消费这个局部向量，最后做一次AllReduce合拢成完整结果。这一步的通信是注意力子层里最大的一笔，也是后面说序列并行时重点优化的对象。

3.2 MLP：两个线性层是一对“先分后合”的搭档

MLP有两个线性层。以常见的SwiGLU类结构为例，第一层把hidden维映射到若干倍的中间维，第二层把中间维映射回hidden维。TP的切法是：第一层用列并行，第二层用行并行。

第一层的列并行：把中间维切成N段，每卡只计算自己那一段，输出是中间维的局部切片。这里有个工程细节：如果你用的是gate和up双分支结构，两个分支矩阵的输出维度都指向中间维，实现时通常会直接把两段拼成一个更大的投影矩阵。无论怎么拼，都必须保证中间维能被N整除，这样每卡拿到的中间块才是完整等份。

第一层的输出经过激活函数后，进入第二层。我在2.3小节说的配对关系在这里正好生效：第二层用行并行，输入是上一层列并行后沿中间维切出的N份。每张卡把自己这份中间维和对应的第二层权重行块相乘，得到hidden维的部分和，最后AllReduce得到完整hidden维输出。一个MLP子层就这样以“列并行进门、行并行出门”的方式结束。

3.3 归一化层：不聚合就统计错，这是最容易忽略的一环

很多第一次自己写TP的人会被LayerNorm绊住。因为TP把hidden维切开之后，LayerNorm要计算整个hidden维上的均值与方差，RMSNorm要计算整个维度的平方均值。如果输入X在卡间只是局部hidden切片，单独统计本卡局部数值，算出来的归一化结果完全是错的。必须先把上一层行并行输出的完整hidden向量拿到本地，再做归一化。

工程上有两种选择：一种是在每个子层的出口把输出始终保持完整，代价是每个TP层之间都有一轮全量聚合，通信非常重；另一种是让归一化层的输入保持按hidden切分，但把统计量的归约做小——例如先对局部统计量做一次归约得到全局均值或方差，再用公式修正。具体实现因框架而异，但底层思路不变：归一化需要知道整个通道维的统计信息，无论TP怎么切，这部分信息都不能省。

RMSNorm因为只需要二阶统计量，在TP下做局部平方和加总，再广播回去，实现很干净。这也是为什么很多LLM普遍喜欢RMSNorm的原因之一——它天然便宜，且更适配TP和序列并行。我曾在某并行框架的源码里翻到过这种实现：对局部张量做平方求和，归约一个很小的统计量，然后把归一化系数广播给所有卡，思路非常漂亮。

4. 通信才是TP的命门：一张图看清集合通信与开销

4.1 行并行和列并行各自要哪种通信算子

通信是TP的灵魂。如果你打开日志发现TP版本比单卡还慢，十有八九是通信开销没有压过计算收益。四个最常碰到的算子是AllGather、AllReduce、ReduceScatter和Broadcast。列并行的输出拼接用AllGather：每卡有一份局部输出，最后要聚合成一个完整张量。反向传播时，梯度的切分回传用ReduceScatter，把完整梯度归约后按位置散到各卡。行并行的输出求和使用AllReduce：每卡有部分和，最终每一张卡都要得到完整结果。反向传播时，输入梯度的归约同样需要AllReduce。

整理成一张表，方便对照：

这种对称性其实揭示了一个规律：TP内部的每一层，输入和输出的数据分布只有两种状态，完整与分片。设计TP网络时，所有算子和切分纬度的选择，本质上是在这两种状态之间做路由。通信算子本身不难，难的是搞清楚每一步张量到底处于哪个状态、下一层期望哪个状态，中间差一个状态就要补一轮通信。

4.2 通信量公式：为什么TP=8并不比TP=1快8倍

我常看到有人以为TP=8就是单卡计算量的1/8，速度自然变8倍。真跑起来，通信会合起来抽走很大一部分收益。TP中每一层前向的通信量大致正比于输入张量大小M＝batch×seq×hidden。列并行的AllGather需要把N份局部块汇总，通信量约为M×(N-1)/N；行并行的AllReduce要归约部分和，通信量同样接近M。一个标准Transformer子层至少包含一次注意力输出投影的AllReduce和一次MLP第二层的AllReduce，中间还有多次AllGather，所以单层前向的通信量通常在2到4倍M的量级，反向还要再翻倍。

于是问题来了：计算量随N线性下降，通信量却随N增加。当N=8时，单卡计算只剩1/8，可通信数据量却是TP=1时的好几倍。只有卡间互联带宽足够大，通信和计算能重叠，才能看到近线性加速。这也是为什么TP几乎只被用在单节点内，靠的还是高性能卡间互联，而不是普通以太网交换机。

为了更直观，我通常用下面这个粗略估算方式：假设某层计算时间在TP=1时是T_compute，通信时间是T_comm；TP=N后计算时间约等于T_compute/N，通信时间约等于C×T_comm（C随N缓慢增长）。只要N增加到某个阈值后，通信时间超过计算缩减的收益，吞吐就会停滞甚至倒退。这个阈值在不同的互联带宽下完全不同，必须用实测数据去定。

4.3 序列并行：把TP里最贵的两次通信省下来

有一个非常重要的优化叫序列并行，它通常和TP绑在一起出现。核心观察是：像LayerNorm、Dropout这类对通道维不敏感的逐元素算子，其实不需要等到完整hidden张量聚合后再执行，按序列维度切分也同样正确。

结合TP后，序列并行的做法是：在每个子层中把LayerNorm和Dropout等算子放在序列维的切片上处理，只有矩阵乘法需要跨卡通信。经典TP中两个行并行输出位置的AllReduce，可以改造成ReduceScatter，把归约结果按序列片段散到各卡；下一层QKV投影需要完整序列输入时，再做AllGather收回来。这样通信规模从“整条序列完整张量归约”变成“每卡只归约自己那一段序列”，激活存储也随序列分片大幅下降。

说人话就是：序列并行把TP通信里最贵的两次全量聚合，切成更小粒度的分片通信，同时连带降低了激活值显存。这也是为什么今天主流大模型训练的“TP”基本都默认绑定序列并行一起用。如果你听到别人说“开了TP”，多半他同时开了序列并行，二者在工程实现里几乎粘在一起。

5. 组合与选型：TP不是孤军奋战，它后面还站着PP、DP和显存优化

5.1 三维并行的分工：TP、PP、DP各管哪一段

在真实的多卡训练中，很少只开TP一种并行，通常会把卡组织成多维结构：数据并行维度、流水线并行维度、张量并行维度。

数据并行最直观。每张卡或每组卡处理不同的batch数据，梯度定期归约。流水线并行则按层切，把模型的若干连续层分配到不同设备，设备间传送层与层之间的激活张量。张量并行负责层内的矩阵切分。三者组合时，原则通常是：优先在节点内使用大TP，因为TP通信最密集；跨节点之间使用流水线并行或数据并行，因为它们的通信频率更低、对带宽要求更宽松。

一个很常见的布局是：一个节点内的8张卡构成一个TP组，节点间跑数据并行或流水线并行。这样TP的通信被限制在节点内部，带宽充足；跨节点的通信只承担DP的梯度归约或PP的层间激活，更适合普通数据中心网络。如果强行让TP跨节点，每秒都在进行的多层聚合会被十字交叉的跨节点通信拖住，整体吞吐会非常难看。

5.2 TP和参数分片的取舍

TP之外，还有一条显存路线：参数分片，在不少主流框架里也常被提起。它把权重、梯度、优化器状态分片到不同卡上，计算时再临时还原。参数分片的优势是不要求数据都在同一节点内，通信以归约为主，跨节点也能工作。缺点是每个Transformer层都要发起通信来还原参数和规约梯度，通信频率和量级并不低，而且不会切到激活值。

TP则刚好相反：通信密集但局域，激活值可以跟hidden维一起切。如果你只是想让一个几十B模型跑起来，参数分片往往更省事；但如果序列很长、激活值爆炸，或者你的卡间互联很快，TP的优势会明显得多。

我自己的经验是：不要让TP和参数分片无序叠加。两者都切权重，组合起来会放大通信，却没有额外显存收益，甚至某些状态仍被重复存储。最常见的工程默认是：TP组内只做TP，把参数分片让给更外层的DP卡组，或者干脆用纯DP+TP的组合。想把所有维度都开一遍，请先做一次通信量估算，别让优化项互相打架。

5.3 显存优化的叠加顺序：TP配重计算和混合精度的经验

TP能解决权重存储问题，但解决不了全部激活值问题。长序列训练下，激活值依然是显存杀手，这时候最常用的组合是激活重计算加混合精度。

我建议的顺序是：先开TP，把模型权重、梯度、优化器切到可接受范围；然后打开混合精度，进一步把存储减半；再做激活重计算，只保留关键中间节点，其他激活用计算换显存；最后才考虑增加TP规模。原因很简单——TP规模每加一档，通信量都会跳一次，边际收益递减明显；而激活重计算不改变通信结构就能释放大量显存，代价只是多算一次前向。

另外，梯度累计、微步长这类调度技巧也能间接缓解显存。它们和TP并不冲突，TP只负责模型存储的切分，激活管理完全交给上层调度。实际训练中，我会先看显存峰值里激活值占多少，再决定要不要开激活重计算；如果激活占比不高，就没必要纯为了减少显存去增加额外的重算开销。

6. 实践中的对齐恐怖症：我从TP配置上踩过的坑

6.1 hidden_size、head数、intermediate_size的整除问题

TP的切分不是随便切，很多维度必须被整除，否则会在启动时就遇到各种难以定位的报错。

第一个是hidden_size，它必须能被TP整除。几十亿参数的模型通常hidden是4096、5120这类数，4或8整除没问题；但如果hidden是某个奇怪的质因数组合，TP=6时就可能崩。第二个是注意力头数。切分注意力时，必须保证每卡拿到整数个完整头，所以head数必须能被TP整除。我曾遇到一个模型head数是20，愣是拿它去跑TP=8，结果头数不均匀，QKV投影的形状直接错乱，报错信息根本看不出是头分配问题。第三个是intermediate_size。MLP中间维度大小必须能被TP整除，尤其是SwiGLU里gate和up矩阵经常拼接在一起，中间维翻倍后的总列数也要对得上。

排查这类问题，最有效的动作是写一个小脚本，在初始化模型之前把TP规模对上述所有维度做一遍求余校验，并在日志里打印每个维度的分片形状。校验不过就立刻退出，这比等训练跑到第三步才随机炸一个shape mismatch要节省无数时间。

6.2 词表大小不整除：Embedding和输出头的隐藏大坑

词表相关的坑尤其隐蔽。Embedding矩阵的形状是(vocab_size, hidden)，输出logits时则是(hidden, vocab_size)。TP对词表维度做列并行时，vocab_size必须能被TP整除。很多预训练词表大小是精心选过的数，但它不一定能被你的TP整除。遇到类似情况，常见做法是把词表padding到能整除的尺寸，再在损失计算时把padding位置的损失过滤掉。这种padding会带来额外嵌入参数和少量计算浪费，但换来了稳定的并行切分。

更大的坑在输出头。如果模型共享embedding与输出投影，词表切分后，logits也要按列并行切，最后算交叉熵。常见错误是仍然在完整logits上算损失，结果把局部logits聚合到完整形状后，显存直接爆掉。正确方法应该是每卡只计算属于自己词表分区的logits与label的部分损失，再用AllReduce归约平均。我最初实现时偷懒直接聚合logits，单卡显存瞬间多出几十GB，那是那次训练里最深刻的教训之一。

6.3 TP大小与卡间互联拓扑：跨节点TP的惨痛教训

我有一次对训练吞吐不满意，以为把TP从4提到8能翻倍，结果实测没提升，反而略降。分析半天，发现我把TP=8的8张卡分布在了两个物理节点上，跨节点传输成了整个系统的瓶颈。TP的AllGather和AllReduce每层必用、频率极高，对节点之间的低带宽链路来说，每步的通信开销几乎吃掉了全部加速收益，甚至倒亏。

正确的做法是把TP严格限制在同一个高速互联域内。如果你只有两个节点、每节点4张卡，那么优先TP=4，节点间用流水线并行或数据并行，而不是强行TP=8。另外，卡到GPU的rank顺序也要对得上拓扑。很多框架初始化时是按进程rank编号分配GPU的，如果TP维度对应的rank跨越了节点边界，就等于把通信引到了慢链路上，光看卡数根本看不出来。

6.4 启动前必查清单：把这些检查完再开始训练

结合这些坑，我整理了一个每次开启TP前都会跑的检查清单：

• 用脚本验证hidden、head数、intermediate、vocab对TP规模的整除性。
• 确认所有TP组落在同一高速互联域内，最好用拓扑感知的启动方式自动绑定到同一交换域。
• 启动一个极小的2层模型做冒烟测试，把TP规模设为目标值，观察各卡shape打印是否符合预期。
• 用性能分析工具记录AllGather、AllReduce耗时占比，如果通信占比超过30%，先检查通信域而不是盲目加卡。
• 跑一个固定随机种子的几步训练，对比TP=1与TP=N的浮点输出是否对齐，提前发现切分逻辑问题。
• 开混合精度和激活重计算后再测一次显存峰值，确认TP切开后的权重存储符合预期。

经历了这些之后，我对TP的态度从“怕它炸”变成了“先查表再启动”。这六项只要都过，后续训练里的TP相关报错基本能少掉一大半。至于更深层的坑，比如某个特定切分算子的实现细节、不同框架对不同维度的处理习惯，那些就需要在真实报错里一点点积累了。