8.3B 的训练状态装不进一张 32GB V100。Megatron 没有把层拆开，而是把每层的矩阵按列、按行各切一刀（attention 与 MLP 两个子块各一套）：
第一刀让每张卡独立算 GeLU，第二刀在出口做一次 all-reduce 把部分和加回来。
代价是每层前向 2 次、反向 2 次 all-reduce；收益是每卡线性层权重和 GEMM 同步缩到 1/TP。
Transformer 的两块计算都能拆成连续两个 GEMM：MLP 是 h → 4h → h，attention 是 QKV → 头内注意力 → 输出投影。
Megatron 把第一个 GEMM 按列切，让每张卡拿到一部分输出通道；把第二个 GEMM 按行切，让每张卡消费自己那份输入。
这样切的关键收益是中间不需要同步：GeLU、softmax 都在本卡的通道/头上独立完成。整个子块只在行并行的出口做一次 all-reduce，把各卡的部分和相加。一个 Transformer 层有 attention 和 MLP 两个子块，所以前向 2 次、反向 2 次。
通信量按激活大小走：每次 b·s·h 个元素，和 TP 度数无关；线性层权重则按 1/TP 缩小。8.3B 模型只用 8 路 TP（8 卡）时弱扩展效率 77%；再叠 64 路 DP、共 512 张 V100 后为 74%（§5.1）。
b·s·h；8.3B 配置（b=8、s=1024、h=3072）fp16 约 48.0 MiB/次先分清两种内存压力：层数堆高了放不下，和单层本身放不下。TP 解决的是后者。
论文的出发点是 8.3B 参数的 GPT-2 式模型（§1）。按 fp16 存权重约 16.6 GB（8.3B × 2 B 的换算，见 LEDGER），而训练时还要放梯度、Adam 的优化器状态和每层的激活；单张 32GB V100（§5）装不下。 一个自然的想法是把 72 层分给 8 张卡，每张卡放 9 层。但这样做的瓶颈不在权重总量，而在每一层自己的 GEMM 和激活：PP 不切单层，单层的 GEMM 和激活尺寸不变；论文因此做层内切分，并与 PP 正交组合。
线性层是 Y = XA。如果按行切 A、按列切 X，得到的是
X₁A₁ + X₂A₂；但 GeLU 是非线性的：
GeLU(X₁A₁ + X₂A₂) ≠ GeLU(X₁A₁) + GeLU(X₂A₂)。
于是必须在 GeLU 之前做一次同步，把部分和加起来——这正好是 TP 想避免的通信点（§3 Eq. 2–3）。
论文的解法是换一个切法：先按列切 A，让每张卡的输出通道互不重叠。这样 GeLU 可以分别在每张卡上做，不需要先求和； 紧接着的第二个 GEMM 按行切，让每张卡用自己的那部分 GeLU 输出继续算，最后只加一次部分和。切分点从「非线性之前」挪到了「整个子块的出口」。
把三条并行轴放在一起看，才能看出 TP 为什么必须存在，以及它为什么总和 DP、PP 一起出现。
每张卡放完整模型，喂不同 batch，反向时 all-reduce 梯度。它能把吞吐做上去，但每卡显存不降：模型参数、优化器状态、激活都还是完整一份。
适合「单卡放得下、想跑得更快」。
把第 1–9 层放卡 0、第 10–18 层放卡 1，以此类推。每卡只存一部分层，但单层的张量仍要完整放得下；并且流水线有气泡，边界还要传激活。
适合「层数很多、单层还算得动」。
还有一条更早的路线是 Mesh-TensorFlow：用一套语言描述张量在哪些维度上分布，再编译出集合通信。论文没有走编译器路线，而是直接在 PyTorch 的 Transformer 实现里插入几个 autograd 算子——
f 在前向是恒等、反向做 all-reduce；g 在前向做 all-reduce、反向是恒等。Transformer 层的切分只靠这两个共轭算子拼出来（§3、Code 1）。
论文明确把 TP 定位成与 PP 正交且互补（§3 末段）。实际训练时常见组合是：TP 放在机箱内（NVLink），PP 跨机箱，DP 放在最外层复制。 8.3B 的配置就是 8 路 TP × 64 路 DP = 512 张卡（Appendix B.1）。
看懂 MLP 的这两刀，attention 的 QKV 和输出投影是同一套动作换个矩阵形状。
第一层是 Y = GeLU(XA)，其中 X 是 [b, s, h]，A 是 [h, 4h]。
把 A 按列切成 [A₁, …, A_TP]，每张卡算自己的 Yᵢ = GeLU(XAᵢ)，输出通道是 4h/TP。
第二层是 Z = YB，B 是 [4h, h]；把 B 按行切成 [B₁; …; B_TP]，每张卡算 Zᵢ = YᵢBᵢ，最后把所有 Zᵢ 相加：
换成 PyTorch 的权重形状（nn.Linear(in, out) 的权重是 [out, in]），列并行切的是输出维：每卡权重 [4h/TP, h]；
行并行切的是输入维：每卡权重 [h, 4h/TP]。两者各自是 4h²/TP 个参数，一层 MLP 合计 8h²/TP。
Multi-head attention 的 Q、K、V 投影可以拼成一个 [3h, h] 的矩阵（GQA 时是 [h + 2hkv, h]）。
论文把 QKV 投影按列切，让每张卡只拿自己负责的那几个头：每卡 h/TP 个头，每个头在自己的卡上完成 softmax(QKᵀ/√d)V，头与头之间不需要通信（§3、Figure 3b）。
随后的输出投影是 [h, h] 矩阵，按行切，出口做一次 all-reduce。
在 Megatron-Core 里，这个融合的 QKV 层是列并行线性层：非 TE 路径为 ColumnParallelLinear（megatron/core/tensor_parallel/layers.py:1107），TE 路径为 TELayerNormColumnParallelLinear（megatron/core/extensions/transformer_engine.py:1551）；attention 模块里的实例化见 megatron/core/transformer/attention.py:1871–1886。
GQA 时权重按 [q₁ q₂ k₁ v₁ | q₃ q₄ k₂ v₂ | …] 交错排列，切分时再拆成 q、k、v 三段（megatron/core/transformer/attention.py:2034–2035、megatron/core/transformer/attention.py:2084–2097）。
输入 embedding 的论文记法是 [h, vocab]，按词表维切分后每张卡存 vocab/TP 列；输入 token 落到别的卡上时置零，最后做一次 all-reduce 得到完整 embedding（§3）。
Megatron-Core 的 VocabParallelEmbedding 权重布局为 [vocab, h]，对应每卡 vocab/TP 行（megatron/core/tensor_parallel/layers.py:351）。
输出 logits 如果先算成 [b, s, vocab] 再 all-gather，通信量是 b·s·vocab，比激活大得多；论文把并行 GEMM 的输出直接接上交叉熵，只传 b·s 个标量损失，把这一块通信压掉了（§3）。Megatron-Core 对应 vocab_parallel_cross_entropy（megatron/core/tensor_parallel/cross_entropy.py:251）。
下面的计算器只算线性层权重，不含 bias、LayerNorm、embedding；通信按 all-reduce 的 payload 算。切换 Column / Row 可以看同一块权重是沿哪个维度被切开。
[f/TP, h]，输出 [b, s, f/TP]；这一层内部不需要通信。
公式：每层参数 (2hf + 4h²)/TP（attention 按 MHA：QKV 3h² + 输出 h²）；单次 all-reduce 载荷 b·s·h 个元素；
每层 4 次（2 前向 + 2 反向）；ring 每卡搬运 8·b·s·h·(TP−1)/TP 个 fp16 元素。
口径：只计线性层权重（不含 bias / LayerNorm / embedding；attention 按 MHA，QKV 3h² + 输出 h²）；权重按 fp16 的 2 B/元素 折算。
通信只计 all-reduce payload；ring 流量采用 2N(TP−1)/TP 的带宽模型，不含 kernel 启动和网络延迟。这是公式计算器，不模拟 NCCL 的实际调度。
A：h=4096，f=16384，TP=8，b=8，s=2048（页面默认值）
每卡每层参数 = (2×4096×16384 + 4×4096²) / 8 = 25,165,824 = 25.17 M
fp16 权重 = 25,165,824 × 2 / 2²⁰ = 48.0 MiB
单次 all-reduce = 8 × 2048 × 4096 × 2 / 2²⁰ = 128.0 MiB
每层 4 次 = 4 × 128.0 = 512.0 MiB
ring 每卡 = 8 × 8 × 2048 × 4096 × 2 × (8−1)/8 / 2²⁰ = 896.0 MiB
B：h=3072，f=12288，TP=8，b=8，s=1024（Table 2 的 8.3B 宽度）
每卡每层参数 = (2×3072×12288 + 4×3072²) / 8 = 14,155,776 = 14.16 M
fp16 权重 = 14,155,776 × 2 / 2²⁰ = 27.0 MiB
单次 all-reduce = 8 × 1024 × 3072 × 2 / 2²⁰ = 48.0 MiB
每层 4 次 = 4 × 48.0 = 192.0 MiB
ring 每卡 = 8 × 8 × 1024 × 3072 × 2 × (8−1)/8 / 2²⁰ = 336.0 MiB
注意 B 组只换了宽度，没有算 72 层、没有算 embedding 和优化器状态；它对应的是「一层 Transformer 的线性层」这笔账。
TP 的通信频率远高于 DP：不是每个 step 一次，而是每个 Transformer 层 4 次。这决定了它通常只放在机箱内的 8 张卡上。
记激活的 hidden 为 h、micro-batch 为 b、序列长为 s、元素字节数为 d。每次 all-reduce 的张量是 [b, s, h]，
每层前向 2 次、反向 2 次，所以每层每卡的 payload 是 4·b·s·h·d 字节；按 ring all-reduce 的带宽模型，每卡实际搬运约 8·b·s·h·d·(TP−1)/TP 字节。
这个量只随 b·s·h 增长，不随 TP 度数下降——TP 降低了每卡的 GEMM，但没有降低每卡要同步的激活。
这就是 TP 与 DP 的物理差别：DP 的梯度 all-reduce 每个 step 一次，张量和参数同大；TP 的 all-reduce 每层 4 次，张量和激活同大。
层数一多，TP 的通信次数是 4 × layers，所以它对延迟和带宽都敏感。论文的 512 卡集群里，机箱内 NVSwitch 是 300 GB/s，跨机箱 InfiniBand 是 100 GB/s（§5），8 路 TP 可以放在一个 DGX-2H 机箱内。
每层 4 次集合通信，张量 [b,s,h]；延迟敏感，放在 NVLink 域内。
PP 只在层边界传激活，DP 每个 step 传一次梯度；都可以跨机箱，常放在 TP 外层。
原始 Megatron-LM 的 TP 只切矩阵，LayerNorm、dropout、残差这些逐元素操作仍在每张卡上完整复制（§3）。Megatron-Core 的 sequence_parallel 把这些区域的激活沿序列维再切一刀：
列并行层在 forward 前做 all-gather 把序列拼回来（megatron/core/tensor_parallel/layers.py:697–698），行并行层在出口做 reduce-scatter（megatron/core/tensor_parallel/layers.py:1761–1764），反向对应互换（megatron/core/tensor_parallel/mappings.py:528–549）。
配置项的注释直接指向 Korthikanti et al. 2022 的 Reducing Activation Recomputation in Large Transformer Models（megatron/core/model_parallel_config.py:123–127）。
结果是：残差流激活从每卡 b·s·h 降到 b·s·h/TP；通信字节不变（4 次 all-reduce 换成 2 次 AG + 2 次 RS 的前向，反向再各两次），但集合通信的调用次数从 4 变成 8。
省下的是显存，不是带宽。
b·s·h，每层 4 次 all-reduce，每卡 ring 流量 896.0 MiB。
另一个 Core 侧的变化是 tp_comm_overlap：把 all-gather / reduce-scatter 与相邻 GEMM 流水起来，用 Transformer Engine 的通信计算重叠能力，而不是等 GEMM 算完再发集合通信（megatron/core/model_parallel_config.py:325–338）。
论文的实现描述只讲到「把通信插进 autograd 图」（§3），没有提通信重叠；这是工程实现层面的差异，不改变上面的切分数学。
源码锚点：Megatron-Core commit 70f7a26。路径均从仓库根写全，例如 megatron/core/tensor_parallel/layers.py。
MLP 和 attention 的矩阵形状不同，但都是「列并行 → 本地非线性 → 行并行 → all-reduce」。
把图读成两句话：列并行负责「每张卡算不同通道」，行并行负责「每张卡消费自己的通道并求和」。 前者消掉了 GeLU 前的同步，后者把同步推迟到子块出口；两者合起来，一层里只剩两次前向 all-reduce（attention 一次、MLP 一次）。
论文给了三组数字：8.3B 模型变大时的弱扩展、1.2B 固定 batch 加卡的强扩展、以及头数对扩展效率的影响。
Table 1 是四组 GPT-2 式配置。hidden size per head 固定 96，靠增加层数和头数把模型从 1.2B 推到 8.3B；8.3B 那行需要 8 路 TP，再叠 64 路 DP 就是 512 张卡。
Table 1（arXiv v4）：扩展实验的配置。最后一列是 TP × DP 的总卡数。
| 模型 | hidden | heads | layers | 参数量 | TP 卡数 | TP + DP 卡数 |
|---|---|---|---|---|---|---|
| 1.2B | 1536 | 16 | 40 | 1.2B | 1 | 64 |
| 2.5B | 1920 | 20 | 54 | 2.5B | 2 | 128 |
| 4.2B | 2304 | 24 | 64 | 4.2B | 4 | 256 |
| 8.3B | 3072 | 32 | 72 | 8.3B | 8 | 512 |
论文报告（§5.1）：8.3B、8 路模型并行达到线性扩展的 77%；叠上 64 路数据并行、总共 512 卡时是 74%，下降 3 个百分点（77 − 74）。 另一个口径来自 Abstract/§1：整个应用在 512 张 V100 上持续 15.1 PFLOPs，相对单卡 1.2B 基线（39 TFLOPs，峰值 30%）的每卡吞吐比值是 76%——它衡量的是「512 卡总吞吐 ÷ 512 ÷ 单卡基线」，不是 §5.1 的扩展效率，两者不要混用。
Table 8 固定 1.2B 模型和 batch=8，只增加 TP 卡数：2 卡 1.64×、4 卡 2.34×、8 卡 2.98×。 4 卡时还有 58.5% 的线性效率，8 卡只剩 37.3%；论文的解释是每卡计算量下降后，内存带宽和通信开销开始主导（Appendix D.2）。 这是「TP 不是越大越好」的直接证据。
Table 7 固定 8.3B 和 8 路 TP，只改注意力头数：16 头（每头 192 维）是 82%，24 头（128 维）是 80%，32 头（96 维）是 77%，从 82% 到 77% 下降 5 个百分点（82 − 77）。 论文的解释是头数增加会让 self-attention 里的 GEMM 变小，同时 softmax 的元素变多（Appendix D.1）。这也是 TP 切分粒度受模型结构影响的例子：切得越碎，单卡 GEMM 的效率越低。
Table 7（arXiv v4）：8.3B、8 路 TP 下，头数对扩展效率的影响。
| Attention heads | Hidden size per head | Scaling efficiency |
|---|---|---|
| 16 | 192 | 82% |
| 24 | 128 | 80% |
| 32 | 96 | 77% |
最终用于语言建模的 8.3B 配置（Table 2 的 8.3B 行；本页只引用这一行）是 72 层、hidden 3072、24 个头、每头 128 维、512 张卡、每 epoch 2.10 天。 它和 Table 1 的扩展配置（32 头）不同；论文只说明了两者不同，没有解释选择原因。 推断：24 头在 Table 7 里比 32 头的扩展效率高 3 个百分点（80 − 77），可能是精度与吞吐之间的折中；Table 7 只报告扩展效率，没有对应的下游精度。
Table 3 是零样本结果。8.3B 的 WikiText103 困惑度是 10.81，对比此前最佳的 15.79，绝对低 4.98（15.79 − 10.81），相对低 31.5%（4.98 / 15.79）。 LAMBADA 准确率 66.51%，对比 63.24%，绝对高 3.27 个百分点（66.51 − 63.24），相对高 5.17%（3.27 / 63.24）。 论文没有做「同参数量、不同 TP 度数」的精度对照，所以不能把精度收益直接归给 TP。推断：结果主要来自模型变大，TP 的作用是让这个尺寸训得动。
Table 3（arXiv v4）：零样本评测。困惑度越低越好，准确率越高越好。
| 模型 | WikiText103 困惑度 ↓ | LAMBADA 准确率 ↑ |
|---|---|---|
| 355M | 19.31 | 45.18% |
| 2.5B | 12.76 | 61.73% |
| 8.3B | 10.81 | 66.51% |
| 此前最佳 | 15.79 | 63.24% |
下面这些条件不满足时，TP 要么直接报错，要么收益被通信吃掉。
ColumnParallelLinear 用 divide(output_size, world_size)（megatron/core/tensor_parallel/layers.py:1213），RowParallelLinear 用 divide(input_size, world_size)（megatron/core/tensor_parallel/layers.py:1608）；
divide 在不能整除时直接断言失败（megatron/core/utils.py:599–608，断言在 ensure_divisibility 599–601）。头数也要能被 TP 整除（megatron/core/transformer/attention.py:365–367）。
每层 4 次 all-reduce，跨机箱的 100 GB/s InfiniBand 很容易把 GEMM 的收益吃掉（§5 的 300 GB/s NVSwitch vs 100 GB/s IB）。常见做法是 TP ≤ 8、放在同一台 DGX 内，PP/DP 往外层放。
num_query_groups < TP 时，每张卡分不到完整的 KV 头；Core 的做法是先 all-gather QKV，再按 rank 切片（megatron/core/transformer/attention.py:2033–2053），多了一次通信，也不再是「头内本地计算」。b·s 很小时，每卡 GEMM 太小，all-reduce 的固定延迟占比更高。sequence_parallel）。实际训练通常三者叠加：TP 在最内层（机箱内），PP 跨机箱切层，DP 在最外层复制。Megatron 的 8.3B 配置就是 8 路 TP × 64 路 DP（Appendix B.1）；PP 是论文明确说「正交且互补」的另一条轴（§3）。