大模型分布式训练与推理中的通信原语：从基础到工程实践

目录

  1. 通信原语基础
  2. 训练阶段通信
  3. 推理阶段通信
  4. 优化与基础设施

一、通信原语基础

1.1 集合通信的分类体系

分布式系统中的通信操作可按数据流向分为两大类：一对多（One-to-Many） 和 多对多（Many-to-Many）。

1.2 四大核心原语详解

AllReduce

作用：将所有节点上的数据进行规约（如求和、取均值），并将结果广播给所有节点。每个节点最终持有完整的规约结果。

通信量：设数据大小为 Φ\Phi，节点数为 NN，AllReduce 通信量为 2Φ⋅N−1N≈2Φ2\Phi\cdot\frac{N-1}{N} \approx 2\Phi（Ring-AllReduce 实现）。

AllGather

作用：每个节点持有数据的一个分片，操作完成后所有节点都拥有全量数据的拼接。不做规约，只做收集。

通信量：Φ⋅N−1N≈Φ\Phi\cdot\frac{N-1}{N} \approx \Phi

ReduceScatter

作用：将各节点的数据先全局规约，再将规约结果均匀分散到各节点。每个节点最终只持有规约结果的一个分片。它是 AllReduce 的"前半段"。

通信量：Φ⋅N−1N≈Φ\Phi\cdot\frac{N-1}{N} \approx \Phi

All-to-All

作用：每个节点向其他每个节点发送不同的数据分片，同时接收来自每个节点的不同分片。可理解为一次全局"矩阵转置"，常用于专家并行（MoE）中的 token 路由。

1.3 AllReduce = ReduceScatter + AllGather

这是理解 ZeRO 优化器的关键等式：

为什么 ZeRO 偏好 ReduceScatter + AllGather？

传统数据并行的 AllReduce 要求每张卡同时持有完整梯度。ZeRO（Zero Redundancy Optimizer）的核心思想是消除冗余：

ZeRO 阶段 分片内容 通信模式
ZeRO-1 优化器状态 ReduceScatter（梯度）+ AllGather（参数更新）
ZeRO-2 优化器状态 + 梯度 ReduceScatter（梯度）+ AllGather（参数更新）
ZeRO-3 优化器状态 + 梯度 + 参数 AllGather（前向参数）+ ReduceScatter（反向梯度）

关键洞察：ReduceScatter + AllGather 与 AllReduce 的总通信量相同（均为 2Φ2\Phi），但 ReduceScatter 允许在通信完成之前就开始用分片更新优化器状态，将大内存峰值打散到时序上，从而在通信量不变的前提下大幅降低显存占用。

二、训练阶段通信

2.1 三种并行策略全景

2.2 数据并行（DP）

原理：每个 GPU 持有完整模型副本，处理不同的数据 mini-batch，反向传播后同步梯度。

通信量估算：

设模型参数量为 Φ\Phi，参数类型为 FP16（2 bytes/param）：

每次迭代通信量=2×2Φ=4Φ bytes\text{每次迭代通信量} = 2 \times 2\Phi = 4\Phi \text{ bytes}

以 LLaMA-70B（Φ≈7×1010\Phi \approx 7\times10^{10}）为例：4×7×1010≈280 GB4 \times 7\times10^{10} \approx 280\text{ GB}/迭代

为什么 AllReduce 是必须的？

SGD 的正确性要求所有副本在参数更新时使用相同的梯度（或其平均值）。若不同步，各副本的参数会发散，等效于多个模型各自训练而非协同。

2.3 张量并行（TP）

张量并行将单个算子（如矩阵乘法）沿特定维度切分到多张 GPU 上。以 Transformer 的两层 MLP 为例：

前向传播通信：

  • 列并行层（W1）：输入 X 需要 AllGather 或事先广播（若 X 已在各卡上）
  • 行并行层（W2）：输出需要 AllReduce（汇聚部分和）

反向传播通信：

  • 梯度计算是前向的镜像，方向相反但通信操作对称：前向是 AllReduce，反向对应处是 AllGather；前向是 AllGather，反向对应处是 ReduceScatter

通信量分析：

设序列长度 SS，batch size BB，hidden size HH，TP 度为 tt：

  • 每次 AllReduce 通信量 ≈2BSH\approx 2BSH
  • 一个 Transformer 层（含 Attention + MLP）前向约需 4 次 AllReduce，反向约 4 次
  • 总计约 8 次 AllReduce，通信量 =8×2BSH= 8 \times 2BSH

为什么 TP 延迟最敏感、通信量最大？

  1. 频率极高：每一层的前向和反向都需要通信，无法 overlap（激活值依赖通信结果才能继续）
  2. 强数据依赖：AllReduce 是同步阻塞操作，延迟直接串在关键路径上
  3. 必须部署于节点内 NVLink：带宽 600+ GB/s，相比跨节点 IB 的 100-400 Gb/s 高出 10-20 倍

2.4 流水线并行（PP）

流水线并行将模型按层切分到不同设备，采用 P2P（点对点）通信，而非集合通信。

流水线气泡（Bubble）

朴素 1F1B 调度中，流水线需要"填充"和"排空"阶段，产生空闲等待：

气泡率公式：

Bubble Rate=p−1m+p−1\text{Bubble Rate} = \frac{p-1}{m + p - 1}

其中 pp 为流水线阶段数，mm 为 micro-batch 数。mm 越大，气泡率越低。

Interleaved 1F1B（交错调度）

将每个 Stage 分配多个不连续的层块，使流水线"交错"运行：

气泡率降为原来的 1v\frac{1}{v}（vv 为每 Stage 的 chunk 数），但通信次数增加了 vv 倍（更多 P2P 传输）。

2.5 三种并行策略对比总结

维度 数据并行（DP） 张量并行（TP） 流水线并行（PP）
通信操作 AllReduce / ReduceScatter+AllGather AllReduce（前后向各层） P2P Send/Recv
通信内容 梯度（FP16/FP32） 激活值中间结果 激活值（前向）、梯度（反向）
通信频率 每次迭代一次 每层前后向各多次 micro-batch 边界
通信模式 集合通信 集合通信（高频同步） P2P 点对点
典型部署 跨节点（IB/RoCE） 节点内（NVLink） 跨节点（IB/RoCE）
主要开销 通信量大但可 overlap 串行阻塞，延迟敏感 流水线气泡

2.6 千亿参数模型梯度同步通信量估算

以 GPT-3 175B 为例：

量级结论：百亿参数级别每次迭代梯度通信约 数十 GB；千亿参数级别约 数百 GB 至 TB 量级，必须借助高速互联和 ZeRO 分片。

三、推理阶段通信

3.1 推理 vs 训练：通信量为何大幅减少

核心原因：推理阶段没有反向传播，因此：

  • 不需要梯度同步（AllReduce），消除了最大的通信开销
  • 不需要保留中间激活值用于反向计算
  • 数据并行在推理时各实例独立，无需跨实例通信

3.2 推理中的张量并行与流水线并行

特征 训练时 推理时
TP AllReduce 次数 每层前后向各多次 每层前向一次（减半）
PP P2P 通信 前向激活值 + 反向梯度 仅前向激活值
batch size 大（提高 GPU 利用率） 通常较小（延迟敏感）
通信 overhead 比例 ~20-30% 总时间 ~5-15% 总时间

3.3 Prefill-Decode 分离架构

现代大模型推理分为两个阶段，它们的计算特征截然不同：

分离架构引入的特有通信：KV Cache 迁移

KV Cache 传输量估算（以 LLaMA-70B，seq_len=2048 为例）：

KV Cache=2×80×2×2048×8×128×2 bytes≈10.5 GB\text{KV Cache} = 2 \times 80 \times 2 \times 2048 \times 8 \times 128 \times 2\text{ bytes} \approx 10.5\text{ GB}

这是 Prefill-Decode 分离架构的主要额外开销，对节点间带宽有较高要求。

3.4 KV Cache 在推理通信中的角色

KV Cache 是推理通信优化的核心战场：

  • 本地 Cache：最优，无额外通信
  • PD 分离场景：一次性传输，后续 Decode 本地命中
  • 分布式 KV（未来方向）：支持超长上下文，但通信开销显著

四、优化与基础设施

4.1 减少推理通信开销的常见手段

通信计算重叠（Overlap）示意：

特性 NVLink（H100 NVSwitch） InfiniBand HDR/NDR RoCEv2
带宽 900 GB/s（双向） ~25-50 GB/s/端口 ~25-50 GB/s/端口
延迟 ~1 μs ~1-2 μs ~2-5 μs
适用场景 节点内 TP 跨节点 DP/PP/EP 跨节点（以太网）
丢包处理 硬件保证无丢包 基于 IB 协议保证 需 PFC+ECN 控制
成本 高（专有互联） 高 中（以太网复用）

4.3 为什么训练对零丢包要求极高

集合通信对丢包的放大效应：

AllReduce 是 同步栅栏（Synchronization Barrier）——只有当所有参与者都完成，结果才可用。1000 张卡中哪怕 1 张因丢包重传延迟了 500ms，其余 999 张都必须等待，有效利用率降为零。

这正是 InfiniBand 和配置了 PFC（Priority Flow Control）+ ECN（Explicit Congestion Notification） 的 RoCEv2 在训练集群中的核心价值——硬件级别的无损传输保证。

4.4 NCCL 及其在分布式训练中的作用

NCCL 的核心职责：

NCCL Ring-AllReduce 工作原理：

总结：通信全景视图

一些计算

显存估算

权重显存

weight_mem≈params×bytesshards\text{weight\_mem} \approx \frac{\text{params} \times \text{bytes}}{\text{shards}}

其中 bytes 取决于精度（FP32=4、BF16/FP16=2、FP8=1、INT4=0.5），shards 为张量并行或 ZeRO 分片数。

单 token KV 显存

kv_per_token≈2×L×n_kv_heads×d_head×bytes\text{kv\_per\_token} \approx 2 \times L \times n\_kv\_heads \times d\_head \times \text{bytes}

  • 系数 22：K 和 V 各一份
  • LL：Transformer 层数
  • n_kv_heads×d_headn\_kv\_heads \times d\_head：每层 KV 头的总维度（GQA/MQA 时 n_kv_heads<n_headsn\_kv\_heads < n\_heads）

总 KV 显存

kv_total≈kv_per_token×alive_tokens\text{kv\_total} \approx \text{kv\_per\_token} \times \text{alive\_tokens}

alive_tokens 是当前驻留显存的 token 总数（含所有并发请求）。

KV Cache 管理

Prefill Chunk 数

num_chunks=⌈prompt_lenchunk_size⌉\text{num\_chunks} = \left\lceil \frac{\text{prompt\_len}}{\text{chunk\_size}} \right\rceil

Chunked Prefill 将长 prompt 切成固定大小的块逐批处理，避免单次 Prefill 独占 GPU。

Paged KV Block 数

num_blocks=⌈TB⌉\text{num\_blocks} = \left\lceil \frac{T}{B} \right\rceil

其中 TT 为序列长度（token 数），BB 为每个 page block 的 block size（slot 数）。

Paged KV 内部碎片

waste=⌈TB⌉×B−T\text{waste} = \left\lceil \frac{T}{B} \right\rceil \times B - T

最坏情况（T mod B=1T \bmod B = 1）时碎片接近 B−1B-1 个 slot；BB 越小碎片越少，但管理开销越大。

Shared Prefix 节省的 Prefill token 数

saved_prefill_tokens≈(N−1)×P\text{saved\_prefill\_tokens} \approx (N - 1) \times P

NN 个请求共享长度为 PP 的公共前缀时，只需 Prefill 一次前缀，其余 N−1N-1 次直接命中 prefix KV Cache。

CUDA Kernel 参数速算

Reduce 轮数

rounds≈log⁡2(blockDim)\text{rounds} \approx \log_2(\text{blockDim})

树形归约（tree reduction）每轮活跃线程减半，总轮数即以 2 为底的对数。

Shared Memory 粗算

smem≈blockDim×cache_per_thread×sizeof(dtype)\text{smem} \approx \text{blockDim} \times \text{cache\_per\_thread} \times \text{sizeof(dtype)}

每个线程在 shared memory 中缓存若干元素，乘以 block 内线程总数和数据类型字节数。

gridDim（网格大小）

grid=⌈NblockDim×items_per_thread⌉\text{grid} = \left\lceil \frac{N}{\text{blockDim} \times \text{items\_per\_thread}} \right\rceil

总元素数 NN 除以每个 block 处理的元素数（= blockDim × 每线程处理量），向上取整得所需 block 数。

通信量估算

Ring ReduceScatter 每卡通信量

≈N−1N×S\approx \frac{N-1}{N} \times S

NN 卡、数据总量 SS。Ring 中每卡发送 N−1N-1 轮，每轮发送 S/NS/N，共 (N−1)/N×S(N-1)/N \times S。

Ring AllGather 每卡通信量

≈N−1N×S\approx \frac{N-1}{N} \times S

与 ReduceScatter 对称：每卡将自己的 S/NS/N 分片广播给其他 N−1N-1 卡，总发送量相同。

Ring AllReduce 每卡通信量

≈2×N−1N×S\approx 2 \times \frac{N-1}{N} \times S

AllReduce = ReduceScatter + AllGather，通信量恰好翻倍。当 NN 较大时趋近于 2S2S。

MoE All-to-All 单向通信量

≈T×k×H×b\approx T \times k \times H \times b

  • TT：本卡 token 数
  • kk：每 token 路由的 expert 数（top-k）
  • HH：hidden size（每个 token 的向量维度）
  • bb：数据类型字节数

每个 token 需要把自身 hidden state 发给 kk 个 expert 所在卡，单向总量如上。

MoE All-to-All 往返通信量

≈2×T×k×H×b\approx 2 \times T \times k \times H \times b

dispatch（发送 token 到 expert）和 combine（收回 expert 输出）各一次，共两倍。

每 expert 平均 token 数

avg_tokens_per_expert≈T×kE\text{avg\_tokens\_per\_expert} \approx \frac{T \times k}{E}

EE 个 expert 均匀分摊所有路由流量时的期望值；实际中 load imbalance 会导致热门 expert 远超此值。

量化压缩比

BF16 → FP8：KV Cache 约减半

精度 位宽 字节数/元素
BF16 16 bit 2 bytes
FP8 8 bit 1 byte

FP8BF16=12⇒KV footprint 约减半\frac{\text{FP8}}{\text{BF16}} = \frac{1}{2} \Rightarrow \text{KV footprint 约减半}

BF16 → INT4：权重约变为 1/4

精度 位宽 字节数/元素
BF16 16 bit 2 bytes
INT4 4 bit 0.5 bytes

INT4BF16=0.52=14⇒权重 footprint 约变为原来的14\frac{\text{INT4}}{\text{BF16}} = \frac{0.5}{2} = \frac{1}{4} \Rightarrow \text{权重 footprint 约变为原来的} \tfrac{1}{4}

推导细节：INT4 用 4 bit 表示一个参数，两个参数打包进 1 byte（pack-2）。BF16 每参数 2 bytes，因此 2÷0.5=42 \div 0.5 = 4，压缩比正好为 4×。实际实现中通常每 128 个 INT4 参数存一组 BF16 scale（group-wise quantization），会带来约 1–2% 的额外开销，粗估时忽略。

关键结论速查

问题 答案
DP 通信量 ≈2Φ\approx 2\Phi bytes/迭代（AllReduce）
TP 通信瓶颈 高频同步 AllReduce，串在关键路径，必须用 NVLink
PP 主要损耗 流水线气泡，Bubble Rate=(p−1)/(m+p−1)\text{Bubble Rate}=(p-1)/(m+p-1)
推理 vs 训练通信 推理无梯度，通信量减少 50%+
PD 分离额外开销 KV Cache 跨节点传输，数 GB 量级
零丢包为何重要 集合通信全局同步，1 张卡延迟 = 全集群等待
ZeRO 核心思想 ReduceScatter 替代 AllReduce，通信量不变但消除冗余副本
NCCL 核心价值 拓扑感知 + GPUDirect RDMA + Ring AllReduce 优化实现