TP 的 communication：一层几次？每次多少数据？
TP 越大越快吗？

英文问法How many all-reduces does tensor parallelism need per layer? · How much data does each one move? · Is a larger TP degree always faster? · How would you optimize the TP all-reduce?

考过阿里（"tp的allreduce的次数"）· 华为机考（"张量并行度为 2 时要做一次 All-Reduce，待同步张量大小为 80MB，网络带宽为 40GB/s…"）· 腾讯混元（"TP 并行规模越大训练速度一定越快吗？"）· 快手（"有什么思路优化中间的allreduce吗？"）

先弄清楚：每个概念是什么

f / g operator
每个 TP block（attention、MLP）入口一个 f、出口一个 g。f 只在 backward all-reduce，g 只在 forward all-reduce，切法见 TP Linear。
Ring all-reduce
数据切成 \(t\) 块，reduce-scatter 和 all-gather 各转 \(t-1\) 步。每卡收发 \(2(t-1)/t \cdot S\)，几乎不随 \(t\) 变小，数字见下面的表，细节见 all-reduce。
NVLink vs InfiniBand
机内 GPU 之间走 NVLink，跨机走 InfiniBand 网卡。例：DGX H100 一台 8 卡，每卡 NVLink 900 GB/s（官方口径，按双向合计理解，单向约 450）。跨机每卡一张 400 Gb/s（= 50 GB/s）的网卡，差约 9 倍。
Critical path
TP 的 all-reduce 夹在两个算子中间，下一个 GEMM 要等它的结果才能开始。DP 的 gradient 算完一层就能先发，前面的层接着算，见 overlap。

看图：4 次 all-reduce 在哪，ring all-reduce 怎么走

GPU0GPU1输入 X：每张卡都有完整的一份f：forward 原样复制，backward 在这里 all-reduce ③QKV (column)：head 1–16QKV (column)：head 17–32attention：本卡算attention：本卡算out proj (row)：partial sumout proj (row)：partial sumg：forward 在这里 all-reduce ①，partial sum 相加residual + LayerNorm：每卡一整份f：forward 原样复制，backward 在这里 all-reduce ④fc1 (column)：前一半列fc1 (column)：后一半列GeLU：本卡算GeLU：本卡算fc2 (row)：partial sumfc2 (row)：partial sumg：forward 在这里 all-reduce ②进入下一层attentionMLP
图 1：t = 2 时一层 transformer 的数据流（Qwen3-8B 的 32 个 head）。实线框是 forward 的 all-reduce（① ②），虚线框是 backward 的（③ ④），每层共 4 次，每次的数据是这一层的输出 [b, s, h]。
GPU0GPU1A 半B 半A 半B 半12341 + 3 = 4234 + 2 = 64646开始各有一半的和都有完整的和第 1 步reduce-scatter第 2 步all-gather送出 B 半：2送出 A 半：3送出 A 半：4送出 B 半：6
图 2：两张卡做一次 all-reduce。第 1 步互相发半份并相加，第 2 步交换加好的半份，每张卡一共发 2 次半份，也就是 1 份 S。卡再多，每卡也不超过 2S。

30 秒回答

Each layer has two TP blocks, attention and MLP. Each needs one all-reduce in forward, after the row-parallel linear, and one in backward, before the column-parallel linear: four per layer per micro-batch. Each all-reduce carries b·s·h elements, 64 MiB in bf16 for b = 1, s = 4096, h = 8192. A ring all-reduce moves 2(t−1)/t of that per GPU, so traffic barely drops as t grows, while compute per GPU falls as 1/t. So bigger TP isn't always faster: communication starts to dominate, and past one node it goes over InfiniBand, an order of magnitude slower. Keep TP within the NVLink node; use PP across nodes.

推导 / 关键公式

\[ \text{次数} = 2_{\,\text{fwd}} + 2_{\,\text{bwd}} = 4\ \text{次 / 层 / micro-batch} \] \[ S = b\,s\,h \times 2\ \text{bytes（bf16）} \] \[ T_{\text{AR}} \approx \frac{2(t-1)}{t}\cdot\frac{S}{B} \]

\(B\) 是每卡带宽（NCCL 的 bus bandwidth）。每层每卡 forward + backward 合计 \(8bsh\,(t-1)/t\) 个元素。

为什么 t 大了会亏（以 forward 为例）

\[ \frac{\text{每卡 FLOPs}}{\text{每卡通信 bytes}} \approx \frac{24bsh^2/t}{2 \times 2bsh \times 2(t-1)/t} = \frac{3h}{t-1} \]

分子是每层 GEMM 的 FLOPs（见 复杂度）除以 \(t\)，分母是 2 次 all-reduce，每次每卡 \(2(t-1)/t \times 2bsh\) bytes。比值和 \(b\)、\(s\) 无关，\(t\) 每翻一倍就少一半多。

手算例子

华为题：\(t = 2\)、80 MB、40 GB/s → \(2 \cdot \tfrac12 \cdot 80 / 40 = 2\) ms。

\(b = 1\)、\(s = 4096\)、\(h = 8192\)、bf16，\(S\) = 64 MiB：

\(t\)每卡每次流量\(3h/(t-1)\)链路
264 MiB24576NVLink
496 MiB8192NVLink
8112 MiB3511NVLink
16120 MiB1638跨机 InfiniBand

H100 dense BF16 ≈ 989 TFLOPS（官网的 1979 是 sparsity 数字）。计算要盖过通信，每字节需配 989 TFLOPS ÷ 450 GB/s ≈ 2200 FLOPs（NVLink，按单方向算）或 ÷ 50 GB/s ≈ 19800（InfiniBand）。\(t = 8\)：3511 > 2200，够；\(t = 16\) 跨机：1638 远小于 19800，通信时间约是计算的 12 倍。

必答点

追问

怎么优化 TP 的 all-reduce？

开 SP，把 all-reduce 拆成 reduce-scatter + all-gather，通信量不变，LayerNorm、dropout 的 activation 再除以 \(t\)。把 GEMM 切块，和 all-gather / reduce-scatter 流水 overlap（Megatron Core tp_comm_overlap，要求开 SP）。Hopper + NVSwitch 上用 NVLS，把 all-reduce 交给 NVSwitch 做。

实际最优的 t 是多少？

Narayanan et al. 在 DGX A100 上实测，\(t\) 等于单机 GPU 数（8）时吞吐最高，再往上跨机就掉。

来源Shoeybi et al. 2019 §3（pp.4–5，Fig. 4：每层 forward、backward 各 2 次 all-reduce）· Narayanan et al. 2021 §1（p.1）、§3.2（p.5，\(8bsh(t-1)/t\)，Takeaway #1）、§5.4.1（p.10）· nccl-tests PERFORMANCE.md（all-reduce bus bandwidth \(= 2(n-1)/n\) × algbw）· NCCL env（NVLS）· DGX H100 User Guide（NVLink 900 GB/s，8 × ConnectX-7 400 Gb/s）· H100 spec · Megatron Core model_parallel_config.py（tp_comm_overlap）