Tensor-Parallelism
在 单层内部 把权重矩阵按列/行切到多个 GPU，每个 GPU 只算局部的 matmul，层边界做一次 AllReduce 把 partial sum 合起来。Megatron-LM 的原始方案：Attention 的 QKV projection 按列切、O projection 按行切；FFN 的上投影列切、下投影行切——正好让每层只需一次 AllReduce。代价是每层通信，所以 TP 通常只在 NVLink 域内用（单机 8 卡），跨节点走 Pipeline-Parallelism 或 Data-Parallel。
核心思想
考虑 FFN 的两个 matmul：Y = GeLU(X W₁) W₂。设 4 卡 TP：
- W₁ 按列切：
W₁ = [W₁⁽¹⁾, W₁⁽²⁾, W₁⁽³⁾, W₁⁽⁴⁾]。每卡算Y_i = GeLU(X W₁⁽ⁱ⁾)，得到 hidden 的不同列，无通信 - W₂ 按行切：
W₂ = [W₂⁽¹⁾; W₂⁽²⁾; W₂⁽³⁾; W₂⁽⁴⁾]。每卡算Z_i = Y_i W₂⁽ⁱ⁾，然后Z = ΣZ_i（AllReduce）
Attention 类似：QKV projection 按 head 分组列切（天然适配 multi-head），O projection 行切。
每个 Transformer block 需要两次 AllReduce（Attention 出口 + FFN 出口）。
为什么只在高带宽域内用
每次 forward + backward 要 4 次 AllReduce × 每层 hidden × batch × seq_len 的数据量。对 70B Llama、4K seq、batch 4 而言，每步 TP 通信量 GB 级。NVLink 带宽 900 GB/s 够用，InfiniBand 100-400 Gbps 就是瓶颈。
所以 3D 并行的典型切法：TP 在机内、PP 跨机器、DP 跨 PP stage。
演进
- Megatron-LM (2019)：原始 TP 提法
- Sequence Parallelism (2022)：LayerNorm / Dropout 的输入按 seq 维度切，补充 TP 的 activation 内存未切问题
- Context Parallelism (2024)：长 context 下把 Q/K/V 按 seq 维度切，FA 要做 ring-style 通信（Context-Parallelism）
- TP + EP 协同：MoE 架构下 expert 的权重 TP 与 Expert-Parallelism 交叉切分
反面 / 替代
- 完全 DP + ZeRO：不切权重，只切 optimizer / grad / param 分片。通信量小 per step 但需要 AllGather 参数，适合 small model or high-BW interconnect
- FSDP：ZeRO-3 的 PyTorch 原生实现，和 TP 正交组合
- Expert Parallelism：MoE 专用，把专家切到不同 GPU，通信模式不一样（AllToAll-heavy）
引用本概念的论文
- AXLearn、veScale-FSDP、HexiScale — 训练框架的并行 runtime
- Chakra、ProTrain、BOUTE、BOOST、HetRL、NEST — 训练调度 / 异构训练
- OptiKit、MixLLM、LayeredPrefill、DistCA — 推理阶段 TP 切分
- ParallelKittens、FarSkip-Collective、EventTensor — TP 通信 kernel 优化
- FP8FlowMoE、NVIDIA-Disagg-Study — MoE / Disagg 里 TP 的角色