Tensor Parallelism：把单层切到多卡

谁该读这一篇？ 模型装不进单卡的工程师；想搞清"TP 怎么切、AllReduce 落在哪"的人。 前置阅读： 00-prerequisites.md §7。 耗时： 35 分钟 学完能： 1. 画出 attention QKV / output、MLP gate/up/down 的 TP 切分； 2. 说出 AllReduce 在 Transformer 一层里出现几次； 3. 解释 KV cache 在 TP 下的分布（按 head 切）； 4. 在源码里找到 parallel_state.py 入口； 5. 选 TP=2/4/8 时的 trade-off（速度 vs 显存 vs 通信代价）。

1. 什么时候需要 TP

70B 模型 fp16 ≈ 140 GB → 单卡 80 GB 装不下，必须 TP=2。
70B 模型 FP8 ≈ 70 GB → 单卡刚好；但要给 KV cache 留空间 → 仍可能 TP=2。
DeepSeek-V3 671B 全部 fp8 ≈ 700 GB → TP×PP × EP 混合。

显存和并行度的关系：

\text{per-card weight} = \frac{\text{model total bytes}}{TP \cdot PP \cdot DP\_replicas}

TP 给单卡省最多显存（除以 TP），但代价是每层 AllReduce。

2. 切分策略

2.1 Attention 层

Q = X @ W_Q     # X: [B*L, hidden] · W_Q: [hidden, num_heads*head_dim]
K = X @ W_K
V = X @ W_V
Attn = softmax(QK^T/sqrt(d)) V    # [B*L, num_heads*head_dim]
out = Attn @ W_O                   # [B*L, hidden]

TP=4 切法：

W_Q, W_K, W_V    按 num_heads 维切 4 份  (Column Parallel)
W_O              按 hidden 维切 4 份     (Row Parallel)

每卡持有 1/4 的 heads。 attention 在每卡独立算（heads 之间无交互），完了在 W_O 的 row-parallel 输出后做一次 AllReduce 拼回完整 hidden。

2.2 MLP 层

gate = X @ W_gate    # column parallel
up   = X @ W_up      # column parallel
h    = act(gate) * up
out  = h @ W_down    # row parallel → AllReduce

同样：column → row → AllReduce。

2.3 一层 Transformer 的 AllReduce 次数

Input → LayerNorm → QKV → Attn → W_O → AllReduce ① → +residual
       → LayerNorm → MLP gate/up → MLP_down → AllReduce ② → +residual

2 次 AllReduce / 层。 80 层模型一次 forward 就是 160 次 AllReduce。这是 TP 的主要通信代价。

3. AllReduce 的实际开销

NVLink-4 / NVSwitch (H100):

AllReduce 5 KB 数据：~5 μs
AllReduce 5 MB 数据：~50 μs
一层 forward (BF16 Llama-3 70B, batch=64 decode)：算 ~0.5 ms, AllReduce ~0.1 ms

→ AllReduce 占 ~20% 时间。TP 越多 AllReduce 越慢（不是绝对的，NVSwitch 有 ring topology）。

PCIe 卡（没 NVLink）：AllReduce 慢 5-10×，TP 不划算。

4. KV Cache 在 TP 下

KV cache 也按 head 切：

完整 KV: [seq, num_kv_heads, head_dim]
TP=4:    每卡持有 num_kv_heads / 4 个 head 的 KV

GQA 模型（如 Llama-3）KV head 数远小于 attention head（如 8 vs 64），切的时候要注意：

num_kv_heads 必须能被 TP 整除（否则有 head 重复）。
实际限制 TP 上限。

如果不能整除：复制 head（每卡冗余存）。SGLang 自动处理。

5. 在 SGLang 启用

python -m sglang.launch_server \
    --model-path meta-llama/Llama-3.1-70B-Instruct \
    --tp 4

启动会拉 4 个 Scheduler 子进程，每个绑定一张 GPU。通信靠 NCCL：

INFO ... Init torch distributed begin.
INFO ... Init NCCL begin.
INFO ... Init NCCL end. world_size=4 rank=0,1,2,3

world_size 是总 rank 数，rank 是本进程的 id。

6. 源码入口

srt/distributed/parallel_state.py：

维护全局 _TP_GROUP、_PP_GROUP 等。
get_tp_group() 返回当前 TP 通信组。

srt/distributed/communication_op.py：

tensor_model_parallel_all_reduce(t) —— 模型代码层调用。

模型层的 Linear 层（layers/linear.py）有两个变体：

ColumnParallelLinear：列切，输出按列分布；
RowParallelLinear：行切，自动在 forward 出口做 AllReduce。

新模型实现时直接用这两个，不必自己写 AllReduce。

7. TP × PP 混合（多节点）

单节点 8 卡跑不下时上多节点，要叠加 Pipeline Parallel：

节点 1 (8 卡, TP=8) → 节点 2 (8 卡, TP=8) → 节点 3 (8 卡, TP=8) → ...
   layer 0-19            layer 20-39          layer 40-59          ...

每节点持有一段 layer，节点之间靠 PP send/recv（跨节点慢，需要 RDMA 或者 NCCL TCP）。

SGLang 启用：

--tp 8 --pp 2
# 假设 16 卡，每节点 8 卡，节点间 PP=2

PP 引入 "pipeline bubble"（前几个 micro-batch 在 stage 上跑完才能填满），TPOT 增加。要配合 --max-num-batched-tokens 设大 micro-batch。

详见 scheduler_pp_mixin.py。

8. 选 TP=2/4/8 的取舍

TP	优势	劣势
1	无通信	模型必须装得下
2	通信最少	显存 / 算力翻倍
4	主流 70B 配置	AllReduce 开销出现
8	跑 100B+	AllReduce 占 30%+ 时间

经验法则：

7B / 13B：TP=1。
30B / 70B（H100）：TP=2 或 4。
405B：TP=8 + PP=2。
DeepSeek-V3 671B：TP=8 + EP=8（MoE 专家分散）。

9. 排障 cheatsheet

现象	排查
启动卡在 "Init NCCL"	NCCL 端口冲突 / 网络阻塞；试 `NCCL_DEBUG=INFO`
启动报 "num_kv_heads % tp != 0"	TP 设大了；减或换模型
启动 OOM	TP 设小了；增
跑得慢 / SM 利用率低	TP 太大，通信瓶颈；减
NCCL hang	GPU 之间网络问题；检查 `nvidia-smi topo -m`

详见 08-production-deployment/05-incident-playbook.md。

10. 小结

TP = "一层内部"按列/行切分，每层 2 次 AllReduce。
显存按 TP 切；通信开销随 TP 增长。
KV cache 按 head 切（GQA 时受限）。
主流配置：7B TP=1, 70B TP=2/4, 405B TP=8+PP=2。
源码在 distributed/，模型用 ColumnParallelLinear / RowParallelLinear。

11. 自检

一层 Transformer 在 TP 下有几次 AllReduce？分别在哪？

答案

**2 次**： (1) **Attention 出口**：W_O（row-parallel）输出后做 AllReduce，把 N 卡各算的 partial output 求和拼成完整 hidden。 (2) **MLP 出口**：W_down（row-parallel）输出后做 AllReduce，同上。 Q/K/V 投影和 MLP gate/up 是 column-parallel，不需 AllReduce（输出按列分布到各卡，后续算子直接吃分布形式）。 80 层 Llama-70B 一次 forward 共 160 次 AllReduce；H100 NVLink 上每次 50-100 μs，累计 8-16 ms（decode 时延的 ~20%）。

GQA 模型为什么对 TP 有整除约束？

答案

GQA（Grouped Query Attention）的 KV head 数远小于 Q head（如 Llama-3 70B：Q=64, KV=8）。TP 切分时 KV 也要按 head 切到各卡。约束：`num_kv_heads % tp_size == 0`，否则有卡分不到完整的 KV head。实际：Llama-3-70B (KV=8) 最大 TP=8；TP=4 → 每卡 2 KV head；TP=2 → 每卡 4 KV head；TP=16 不行（8 % 16 ≠ 0）。绕法：**复制 KV head**——每卡持有冗余 KV（每张卡分到自己的 head + 复制别人的），SGLang 自动处理但显存代价大；实战很少这么做，宁愿把 TP 降到 8 + 加 DP 副本。

AllReduce 在 NVLink vs PCIe 的差异多大？

答案

NVLink-4 (H100, 600 GB/s pair): AllReduce 5 MB ~50 μs，160 次 ≈ 8 ms。 PCIe Gen5 (~64 GB/s): AllReduce 5 MB ~500 μs（10×慢），160 次 ≈ 80 ms。Decode 一步本来 20-30 ms，AllReduce 占 80 ms = TPOT 翻 4 倍。所以**没 NVLink 不要做 TP**——单卡装得下就 DP 副本，单卡装不下就接受跨节点 PP 而不是 PCIe TP。 NVLink 拓扑：用 `nvidia-smi topo -m` 看，理想是 NV12（全互联），看到 PIX / NODE 表示 PCIe 经过 PCH，AllReduce 慢。

TP 和 PP 各自的 sweet spot 在哪里？

答案

**TP**：单节点内（NVLink 域内）。8 卡 H100 SXM 是 sweet spot，AllReduce 快。**降时延**（同步通信），但显存切到极致。 **PP**：跨节点。把不同层放不同节点，节点间走 IB / RDMA（带宽窄但延迟 OK）；不需要 NVLink 跨节点。**降时延效果差**（流水气泡），但能跑超大模型。组合：8 卡节点内 TP + 跨节点 PP（如 405B：每节点 8 GPU TP=8，2 节点 PP=2，共 16 GPU）。不要 cross-node TP！跨节点 AllReduce 比节点内慢 10-100×。

跑 405B 模型在 16 卡 (2 节点 × 8 H100) 上，怎么配？

答案

FP16 = 810 GB；fp8 = 405 GB；目标装下 + 留 KV pool 空间。推荐 **TP=8 + PP=2** + FP8 量化： - 每节点 8 卡 TP，节点内 AllReduce（NVLink，快） - 节点间 PP=2，每节点持有一半层（≈202 GB，单节点装下） - FP8 让单节点 80GB×8 = 640 GB 充裕 ```bash # 节点 1 python -m sglang.launch_server --tp 8 --pp 2 --quantization fp8 \ --nnodes 2 --node-rank 0 --dist-init-addr :5000 # 节点 2 python -m sglang.launch_server --tp 8 --pp 2 --quantization fp8 \ --nnodes 2 --node-rank 1 --dist-init-addr :5000 ``` 不推荐 TP=16（跨节点 AllReduce 慢）；不推荐 PP=4 TP=4（PP 气泡多）。

12. 下一步

02-data-parallelism.md — DP 与 router 协同。
03-expert-parallelism.md — MoE EP 切分。
04-disaggregated.md — P/D 分离。
源码：srt/distributed/、srt/layers/linear.py。

上游源码：sglang/python/sglang/srt/distributed/。