张量并行 TP 与列并行 / 行并行的实现

7.1 为什么需要张量并行

看一个具体数字：Llama-3 70B 的一层 FFN 包含 $W_1 \in \mathbb{R}^{8192 \times 28672}$，仅这一个矩阵 fp16 就是 $8192\times 28672\times 2 = 470$ MB；fp32 master + Adam $m,v$ 加起来 $\times 5 = 2.35$ GB。整层带来的状态接近 5 GB。70B 模型 80 层，纯 DDP 是绝无可能。

解决办法之一是把这个大矩阵沿一个维度切到不同 GPU 上，每张卡只算自己负责的"竖条"或"横条"，最后 all-reduce 拼起来。这就是 Megatron 的张量并行（Tensor Parallelism, TP）。

7.2 列并行 vs 行并行：两种切法

Megatron 把所有 nn.Linear 替换成两种特殊版本：

列并行（ColumnParallelLinear）

设输入 $X \in \mathbb{R}^{B \times H}$，权重 $W \in \mathbb{R}^{H \times K}$，普通线性是 $Y = XW$。

把 $W$ 沿列切成 $W = [W_1, W_2, \dots, W_T]$，每个 $W_i \in \mathbb{R}^{H \times K/T}$ 放在 GPU $i$ 上：

$$ Y_i = X W_i,\qquad Y = [Y_1, Y_2, \dots, Y_T] $$

每张卡得到 $Y$ 的一片，不需要通信。反向时需要 all-reduce 输入梯度 $\partial L/\partial X$。

行并行（RowParallelLinear）

把 $W$ 沿行切成 $W = \begin{bmatrix}W_1\\W_2\\\vdots\\W_T\end{bmatrix}$，每个 $W_i \in \mathbb{R}^{H/T \times K}$。输入 $X = [X_1, X_2, \dots, X_T]$（沿 hidden 维切）：

$$ Y_i = X_i W_i,\qquad Y = \sum_{i=1}^{T} Y_i $$

每张卡算出部分和，前向需要 all-reduce。反向时输入梯度不需要通信。

7.3 Megatron 怎么"两次切"组合成 attention / FFN

Megatron 的核心 trick：把 attention / FFN 拆成 "Col → Row" 的组合，把两次通信合并成一次 all-reduce。

FFN（两层 Linear + 激活）

整个 FFN 只在末尾做一次 all-reduce，不是两次。这就是论文里说的 "MLP requires only one all-reduce per layer"。

Attention（QKV + Output Projection）

同样的思想：

$$ \underbrace{\mathrm{QKV} = X \cdot W_{qkv}}_{\text{ColumnParallel}} \;\rightarrow\; \mathrm{Attn}(Q,K,V) \;\rightarrow\; \underbrace{\mathrm{Out} = Y \cdot W_o}_{\text{RowParallel + all-reduce}} $$

QKV 矩阵被合并成一个 column parallel linear，attention 内部每张卡算自己负责的 heads，最后 output projection 是行并行 + all-reduce 拼起来。

7.4 通信量计算

每个 transformer block 的 TP 通信量：

$$ C_\text{block} = 2 \cdot (B \cdot S \cdot H) \cdot \text{bytes} \quad (\text{2 个 all-reduce: attention + FFN}) $$

$B$ batch, $S$ seq length, $H$ hidden size。Llama-70B 单层每张卡每步通信量约 $2 \cdot 4 \cdot 4096 \cdot 8192 \cdot 2\,\text{bytes} \approx 540$ MB（bf16）。整 80 层就是 43 GB/step。这就是为什么 TP 必须配 NVLink——PCIe 跨节点带宽不够，性能会塌方。

7.5 显存收益

设 TP 度 $T$，单个 transformer block 参数从 $P_\text{block}$ 降到 $P_\text{block}/T$。但要注意：

• embedding / norm 不切（除非开 SP）；
• 激活内存只在 attention/FFN 中部省 $1/T$，其他位置（如 LayerNorm 输入）仍然是完整的；
• 这就是为什么后来加了 Sequence Parallel（第 9 章），把 norm / dropout 的激活也切了。

7.6 关键源码地图

7.7 mappings.py 里的 4 个通信原语

这是整个 TP 的"灵魂 4 函数"，每个都有自己的前向 / 反向定义，因为反向时通信方向反过来：

读 Megatron TP 源码的正确顺序：

1. 先读 mappings.py，把 4 个原语吃透；
2. 再读 layers.py，看 ColumnParallelLinear / RowParallelLinear 怎么调这 4 个；
3. 最后看模型层（如 megatron/core/models/gpt/gpt_layer_specs.py），看怎么把 ColumnParallel / RowParallel 拼成 attention + FFN。

7.8 vocab parallel cross-entropy

输出层 logits 维度 = vocab_size，对 Llama-3 是 128K。完整 logits 矩阵在每张卡上是 $B \cdot S \cdot 128000 \cdot 4\,\text{bytes}$，很容易 OOM。Megatron 的招：

1. 把 output embedding 沿 vocab 维切到不同 TP rank（每卡持有 $V/T$ 个 vocab）；
2. 每卡只算自己负责那段 vocab 的 logits 与 loss；
3. 用 cross_entropy.py 里特殊的 VocabParallelCrossEntropy，跨卡 max 找全局 max（数值稳定），跨卡 sum 找 exp sum，然后算 NLL。

这一步把 output layer 的激活内存也降了 $T$ 倍。千亿模型 + 大词表必须开这个。Megatron 中默认就走这条路径。

7.9 dropout 一致性问题

TP 把 attention 切到 $T$ 张卡上，但每张卡都要做同样的 dropout —— 如果 mask 不一致，反向梯度就错。tensor_parallel/random.py 维护两套 RNG：

• data parallel RNG：所有 TP rank 内同步（用于 dataloader、不被切的位置）；
• model parallel RNG：每个 TP rank 不同（用于切了的位置上的独立 dropout）。

这是 TP 实现里最容易踩坑的地方，新人改 dropout 行为时经常忘记切 RNG context。

7.10 调 TP 的经验法则

7.11 与第三方框架的关系

7.12 源码定位：四个核心 TP 类

所有 TP 实现集中在 megatron/core/tensor_parallel/layers.py（54KB），打开它直奔以下四个类：

四个原语在 mappings.py：copy_to / gather_from / reduce_from / scatter_to_tensor_model_parallel_region。配合 autograd 反向自动反演（copy 反 → reduce, gather 反 → scatter, reduce 反 → copy, scatter 反 → gather）。

7.13 MLA：Multi-head Latent Attention（DeepSeek 风格）

2025-Q3 加入。megatron/core/transformer/multi_latent_attention.py:1437 实现，arguments.py:3150 注册 CLI。要点：

• 把 K/V 投影通过低秩 latent 压缩：先到小维度 kv_lora_rank，再恢复，相当于 QKV 上挂了 LoRA；
• 显存收益主要在 KV cache（推理）；训练时计算量略升；
• CLI：--multi-latent-attention --kv-lora-rank 512 --qk-nope-head-dim 128 --qk-rope-head-dim 64 --v-head-dim 128；
• 与 GQA 不冲突，但常单选一个；
• fused 算子在 core/fusions/fused_mla_yarn_rope_apply.py（含 YaRN RoPE 联合融合）。

7.14 TransformerLayer：TP 怎么进到一个 layer 里

看 megatron/core/transformer/transformer_layer.py:291 的 TransformerLayer：

• 继承 GraphableMegatronModule + BaseTransformerLayer（前者支持 CUDA Graph）；
• __init__（line 298）按 TransformerLayerSubmodules（line 232）这一 dataclass 组装：input_norm / self_attn / post_attn_norm / mlp / post_mlp_norm；
• 每个子组件可以替换成 TE 版本（core/extensions/transformer_engine.py）或 MCore 版本，spec 系统就是干这个；
• forward（line 733）按 "norm → attn → residual → norm → mlp → residual" 走，SP 在 norm 前自动 all-gather、之后 reduce-scatter。

7.15 这章你需要带走什么

• 知道 TP 把 weight 矩阵沿列 / 行切到不同 GPU；
• "Column → Row" 组合让 FFN / attention 只需一次 all-reduce；
• num_heads 必须能被 TP 整除；
• mappings.py 的 4 个原语是 TP 的灵魂；
• vocab parallel cross-entropy 处理大词表 OOM；
• TP 不应该跨节点（NVLink-only）。