GPU 显存模型与分布式通信基础

Motivation

如果你问一个 LLM Infra 工程师"大模型系统优化的本质是什么"，大概率会得到两个词：显存和通信。

一个 70B 参数的模型，FP16 下仅参数就占约 140 GB——远超任何单张 GPU 的显存容量。即便模型能塞进一张卡，训练时的梯度、优化器状态、激活值会让显存需求再膨胀 4-8 倍。于是我们不得不把模型"拆"到多张 GPU 上，而"拆"就意味着通信——卡与卡之间需要频繁交换梯度、参数、激活值。

显存决定了能不能跑，通信决定了跑得快不快。 这两个问题贯穿了本系列后续所有文章：数据并行要同步梯度（all_reduce），模型并行要交换激活值（all_gather / reduce_scatter），流水线并行要点对点传递中间结果（send / recv），专家并行要全对全分发 token（all_to_all）。

本文是整个系列的第一篇，我们将：

1. Part A：深入 GPU 硬件架构，理解从寄存器到 HBM 的内存层级，以及为什么 HBM 带宽是 LLM 的核心瓶颈
2. Part B：系统梳理 NCCL 的 8 种通信原语、3 类通信算法、4 种硬件拓扑

掌握这两个基石之后，后续文章中的各种并行策略就不再是"背公式"，而是自然推导的结论。

Part A：GPU 显存模型

GPU vs CPU：不同的设计哲学

在深入 GPU 内存层级之前，先理解 GPU 和 CPU 在设计上的根本区别。

flowchart LR
  subgraph CPU["CPU 设计哲学：延迟优先"]
    direction TB
    C0["强核心 0\n大 Cache | 乱序执行"] ~~~ C1["强核心 1\n大 Cache | 分支预测"]
    C2["强核心 2\n大 Cache | 乱序执行"] ~~~ C3["强核心 3\n大 Cache | 分支预测"]
    CINFO["核心数 8~128\n重点：单线程性能"]
  end
  subgraph GPU["GPU 设计哲学：吞吐优先"]
    direction TB
    S0["SM"] ~~~ S1["SM"] ~~~ S2["SM"] ~~~ S3["SM"]
    S4["SM"] ~~~ S5["SM"] ~~~ S6["SM"] ~~~ S7["SM ..."]
    GINFO["108 SM × 64 CUDA Core\n重点：大规模并行"]
  end

CPU 追求延迟最优（Latency-oriented）：少量强大的核心，配备大容量缓存、乱序执行引擎和复杂的分支预测器，目标是让单条指令流尽快执行完。

GPU 追求吞吐最优（Throughput-oriented）：大量简单的核心，放弃复杂的控制逻辑，把晶体管预算全部花在算术逻辑单元（ALU）上，目标是在单位时间内完成尽可能多的浮点运算。

这就是为什么深度学习天然适合 GPU——矩阵乘法本质上是大规模的、相互独立的乘加运算，完美契合 GPU 的设计哲学。

SM、Warp 与执行模型

GPU 的基本计算单元是 SM（Streaming Multiprocessor）。以 NVIDIA A100 为例，一块 GPU 包含 108 个 SM，每个 SM 内部结构如下：

flowchart TD
  subgraph SM["SM (Streaming Multiprocessor)"]
    direction TB
    subgraph PB["4 个处理块 (Processing Block)"]
      direction LR
      subgraph PB0["处理块 0"]
        direction TB
        P0A["16 FP32 Core\n8 FP64 Core\n1 Tensor Core"]
        P0B["Warp Scheduler\nRegister File"]
      end
      subgraph PB1["处理块 1"]
        direction TB
        P1A["16 FP32 Core\n8 FP64 Core\n1 Tensor Core"]
        P1B["Warp Scheduler\nRegister File"]
      end
      subgraph PB2["处理块 2"]
        direction TB
        P2A["16 FP32 Core\n8 FP64 Core\n1 Tensor Core"]
        P2B["Warp Scheduler\nRegister File"]
      end
      subgraph PB3["处理块 3"]
        direction TB
        P3A["16 FP32 Core\n8 FP64 Core\n1 Tensor Core"]
        P3B["Warp Scheduler\nRegister File"]
      end
    end
    SMEM["Shared Memory / L1 Cache (192 KB, 可配置比例)"]
  end
  PB --> SMEM

Warp 是 GPU 执行的基本调度单位，由 32 个线程组成。同一个 Warp 中的 32 个线程在同一时钟周期执行相同的指令（SIMT，Single Instruction Multiple Threads）。这意味着：

• 如果 Warp 内出现分支（if-else），两个分支会被串行执行（warp divergence），性能损失严重
• 内存访问时，如果 32 个线程访问的地址连续（coalesced access），可以合并成一次内存事务；否则需要多次访问

Occupancy（占用率） 是另一个关键概念。每个 SM 可以同时驻留多个 Warp，当一个 Warp 在等待内存数据返回时（延迟可达数百个时钟周期），SM 的 Warp 调度器会切换到另一个就绪的 Warp 继续执行。这种 延迟隐藏（Latency Hiding） 机制是 GPU 高吞吐的关键——但前提是有足够多的活跃 Warp。占用率越高，延迟隐藏越充分，SM 利用率越高。

内存层级：从寄存器到 HBM

GPU 的内存层级从快到慢依次为：

flowchart TD
  REG["<b>Register File</b>\n~256 KB | ~20 TB/s | 0 cycles\n作用域: per thread"]
  SMEM["<b>Shared Memory / L1 Cache</b>\n192 KB per SM | ~19 TB/s | 1-2 cycles\n作用域: per SM"]
  L2["<b>L2 Cache</b>\n40 MB (A100) | ~5 TB/s | 20-30 cycles\n作用域: 全局"]
  HBM["<b>HBM (Global Memory)</b>\n80 GB (A100) | 2.0 TB/s | 200-400 cycles\n作用域: 全局"]

  REG -->|"~10x 带宽下降"| SMEM -->|"~4x"| L2 -->|"~2.5x"| HBM

  style REG fill:#2d6a4f,color:#fff
  style SMEM fill:#40916c,color:#fff
  style L2 fill:#74c69d,color:#000
  style HBM fill:#b7e4c7,color:#000

逐层说明：

寄存器（Register File）：每个线程私有，访问零延迟。A100 每个 SM 有 256 KB 寄存器文件，分配给 SM 上所有活跃线程。寄存器是最宝贵的资源——每个线程使用的寄存器越多，SM 上能同时驻留的 Warp 就越少，占用率就越低。

共享内存（Shared Memory）/ L1 缓存：SM 内所有线程共享，A100 上每个 SM 有 192 KB，可以在 Shared Memory 和 L1 Cache 之间灵活配置比例。共享内存是程序员显式管理的"软件缓存"，在矩阵乘法的 Tiling 优化中是核心工具——先把数据块从 HBM 搬到 Shared Memory，然后在 Shared Memory 上做多次计算，摊薄 HBM 访问开销。

L2 缓存：全局共享，A100 上有 40 MB。对程序员不可直接控制，由硬件自动管理。

HBM（High Bandwidth Memory）：这就是我们通常说的"显存"。A100 SXM 版本有 80 GB、带宽 2.0 TB/s。虽然 2 TB/s 的带宽看起来已经很高，但对比寄存器的 ~20 TB/s，存在一个量级的差距。任何频繁访问 HBM 的操作都可能成为瓶颈。

一个直观的类比：寄存器像你桌面上随手可拿的便签，Shared Memory 像抽屉里的文件夹，L2 像同一层楼的文件柜，HBM 像隔壁楼的仓库。你肯定希望把最常用的数据放在桌面上，而不是每次都跑去仓库取。

为什么 HBM 带宽是 LLM 的核心瓶颈

要判断一个操作是"算力瓶颈"还是"带宽瓶颈"，核心工具是 Roofline 模型和算术强度（Arithmetic Intensity）。

算术强度 = 计算量（FLOPs） / 数据访问量（Bytes）

以 A100 SXM 为例：

• 峰值算力：312 TFLOPS（FP16 Tensor Core）
• HBM 带宽：2.0 TB/s
• 平衡点：312 / 2.0 = 156 FLOPs/Byte

也就是说，每从 HBM 读取 1 字节数据，至少需要做 156 次浮点运算，才能让计算单元不闲着。低于这个比值的操作就是 Memory-bound（带宽瓶颈），高于则是 Compute-bound（算力瓶颈）。

flowchart LR
  subgraph Roofline["Roofline Model — A100 SXM"]
    direction LR
    MB["🔵 **Memory Bound**\nAI < 156 FLOPs/Byte\n性能 = AI × 2.0 TB/s"]
    BP["⚖️ **平衡点**\n156 FLOPs/Byte"]
    CB["🔴 **Compute Bound**\nAI ≥ 156 FLOPs/Byte\n性能 → 312 TFLOPS"]
    MB --- BP --- CB
  end
  style MB fill:#fff3cd,stroke:#856404
  style BP fill:#cce5ff,stroke:#004085
  style CB fill:#d4edda,stroke:#155724

现在来看 LLM 中的关键操作：

矩阵乘法（GEMM）：对于大尺寸的矩阵乘法 $C = A \times B$，其中 $A \in \mathbb{R}^{M \times K}$，$B \in \mathbb{R}^{K \times N}$，计算量为 $2MKN$ FLOPs，数据量约为 $2(MK + KN + MN)$ Bytes（FP16）。当 $M$、$K$、$N$ 足够大时，算术强度可以很高，是 Compute-bound 的。这就是为什么训练时的 forward pass（大 batch、大矩阵乘法）通常可以充分利用 GPU 算力。

逐元素操作（LayerNorm、GELU、Softmax 等）：每个元素只做几次运算，但都要从 HBM 读一次、写一次。算术强度极低（通常 < 10 FLOPs/Byte），严重 Memory-bound。这就是 FlashAttention 和 kernel fusion 优化的动机——减少对 HBM 的访问次数。

自回归推理（Autoregressive Decoding）：这是 LLM 推理的核心瓶颈。每个 decode step 只生成一个 token，对应的矩阵乘法退化为矩阵-向量乘法（GEMV）：$y = W \cdot x$，其中 $W \in \mathbb{R}^{d \times d}$，$x \in \mathbb{R}^{d \times 1}$。计算量为 $2d^2$ FLOPs，但需要读取整个权重矩阵 $2d^2$ Bytes（FP16），算术强度 = $2d^2 / 2d^2$ = 1 FLOPs/Byte。这比平衡点 156 低了两个数量级，意味着 GPU 算力利用率不到 1%，几乎全部时间在等 HBM 传数据。

这就是为什么 LLM 推理优化中如此强调 KV Cache（减少重复计算的内存访问）、量化（减少权重的字节数）、Continuous Batching（增大 batch size 提高算术强度）。

实战：GPU 显存 Profiling

理论之后我们来看实际数据。以下代码分析了一个 Transformer 模型在单卡上的显存占用分布：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41

import torch
import torch.nn as nn

def profile_memory(model, batch_size, seq_len, vocab_size):
    """分析模型训练时的显存占用"""
    device = torch.device("cuda")
    model = model.to(device)

    # 1. 模型参数
    param_mem = sum(p.numel() * p.element_size() for p in model.parameters())
    print(f"模型参数显存: {param_mem / 1024**3:.2f} GB")

    # 2. 优化器状态 (Adam: 每个参数额外 2 份 FP32 拷贝)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

    # 3. 前向传播，观察激活值
    torch.cuda.reset_peak_memory_stats()
    input_ids = torch.randint(0, vocab_size, (batch_size, seq_len), device=device)

    with torch.cuda.amp.autocast(dtype=torch.float16):
        output = model(input_ids)
        loss = output.sum()

    peak_after_forward = torch.cuda.max_memory_allocated()
    print(f"前向后峰值显存: {peak_after_forward / 1024**3:.2f} GB")
    activation_mem = peak_after_forward - param_mem
    print(f"  其中激活值约: {activation_mem / 1024**3:.2f} GB")

    # 4. 反向传播
    loss.backward()
    peak_after_backward = torch.cuda.max_memory_allocated()
    grad_mem = peak_after_backward - peak_after_forward
    print(f"梯度显存增量: {grad_mem / 1024**3:.2f} GB")

    # 5. 优化器 step
    optimizer.step()
    peak_after_optim = torch.cuda.max_memory_allocated()
    optim_mem = peak_after_optim - peak_after_backward
    print(f"优化器状态增量: {optim_mem / 1024**3:.2f} GB")

    print(f"\n总峰值显存: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")

运行结果大致如下（以一个 1.3B 模型为例）：

模型参数显存:     2.60 GB   (FP16)
前向后峰值显存:   8.50 GB
  其中激活值约:   5.90 GB   ← 激活值是大头！
梯度显存增量:     2.60 GB   (与参数等大)
优化器状态增量:  10.40 GB   (Adam: 2 x FP32 = 4x 参数大小)

总峰值显存:      18.50 GB

这里体现了一个重要的经验法则——Adam 优化器的 4 倍法则：

对于一个 $\Phi$ 参数的模型，仅参数相关的显存就需要 $16\Phi$ bytes，这还不包括激活值。以 70B 模型为例：$16 \times 70 \times 10^9 = 1120$ GB——至少需要 14 张 A100 80GB 才能放下，而且还没有为激活值留余量。

完整代码见 code/01-gpu-memory-distributed/memory_profiling.py

Part B：分布式通信

为什么需要多卡

上一节的计算已经给出了答案：单卡装不下大模型。更具体地说：

当模型分布到多张 GPU 上后，各种并行策略（DDP、FSDP、TP、PP、EP）在不同阶段需要不同类型的通信。NCCL（NVIDIA Collective Communications Library）提供了这些通信操作的高效实现。

NCCL 通信原语：从训练场景推导

理解通信原语最好的方式，不是逐个背定义，而是从实际的训练/推理场景出发，看每个场景自然需要什么样的数据搬运。NCCL（NVIDIA Collective Communications Library）提供了 8 种通信原语，每一种都对应着真实的工程需求。

下面我们假设有 4 张 GPU（Rank 0-3, $P = 4$），数据大小为 $N$，从 5 个场景出发逐一推导。

场景一：DDP — 每张卡算不同数据，梯度怎么同步？

DDP（Distributed Data Parallel） 是最基本的数据并行策略：每张卡持有完整的模型副本，各自处理不同的 mini-batch，然后同步梯度使参数更新一致。

问题很清楚：每张卡算出了自己的局部梯度 $g_i$，我们需要让每张卡都拿到 $\bar{g} = \frac{1}{P}\sum_i g_i$。这正是 All-Reduce 的定义。

All-Reduce

功能：所有卡上的数据做归约（通常是求和），结果存到每张卡。

flowchart LR
  subgraph Before[" "]
    direction TB
    R0B["Rank 0: a0 a1 a2 a3"]
    R1B["Rank 1: b0 b1 b2 b3"]
    R2B["Rank 2: c0 c1 c2 c3"]
    R3B["Rank 3: d0 d1 d2 d3"]
  end
  Before -->|"all_reduce(sum)"| After
  subgraph After["Σi = ai + bi + ci + di"]
    direction TB
    R0A["Rank 0: Σ0 Σ1 Σ2 Σ3"]
    R1A["Rank 1: Σ0 Σ1 Σ2 Σ3"]
    R2A["Rank 2: Σ0 Σ1 Σ2 Σ3"]
    R3A["Rank 3: Σ0 Σ1 Σ2 Σ3"]
  end

通信量：每张卡发送和接收约 $2N \cdot \frac{P-1}{P}$ 数据（Ring 算法）。当 $P$ 很大时趋近 $2N$——与 GPU 数量几乎无关，这是 Ring 算法的精妙之处。

场景二：FSDP — 参数都切碎了，怎么算前向/反向？

DDP 的问题是每张卡都存完整参数——70B 模型根本放不下。FSDP（Fully Sharded Data Parallel） 的思路是：参数、梯度、优化器状态全部按卡切分，每张卡只存 $1/P$。

但切碎之后要计算怎么办？

• 前向传播：计算某一层时，需要该层的完整参数。每张卡只有一个分片，必须临时把所有分片拼起来 → All-Gather
• 反向传播：每张卡算出完整梯度后，需要把梯度归约并重新切分，每张卡只保留自己负责的那份 → Reduce-Scatter

All-Gather

功能：每张卡贡献自己的一个分片，拼出完整数据到所有卡。

flowchart LR
  subgraph Before[" "]
    direction TB
    R0B["Rank 0: a0 _ _ _"]
    R1B["Rank 1: _ b1 _ _"]
    R2B["Rank 2: _ _ c2 _"]
    R3B["Rank 3: _ _ _ d3"]
  end
  Before -->|"all_gather"| After
  subgraph After[" "]
    direction TB
    R0A["Rank 0: a0 b1 c2 d3"]
    R1A["Rank 1: a0 b1 c2 d3"]
    R2A["Rank 2: a0 b1 c2 d3"]
    R3A["Rank 3: a0 b1 c2 d3"]
  end

通信量：每张卡接收 $N \cdot \frac{P-1}{P}$ 数据。

Reduce-Scatter

功能：先归约（求和），再将结果的不同部分分散到不同卡。

flowchart LR
  subgraph Before[" "]
    direction TB
    R0B["Rank 0: a0 a1 a2 a3"]
    R1B["Rank 1: b0 b1 b2 b3"]
    R2B["Rank 2: c0 c1 c2 c3"]
    R3B["Rank 3: d0 d1 d2 d3"]
  end
  Before -->|"reduce_scatter(sum)"| After
  subgraph After["Σi = ai + bi + ci + di"]
    direction TB
    R0A["Rank 0: Σ0 _ _ _"]
    R1A["Rank 1: _ Σ1 _ _"]
    R2A["Rank 2: _ _ Σ2 _"]
    R3A["Rank 3: _ _ _ Σ3"]
  end

通信量：每张卡发送 $N \cdot \frac{P-1}{P}$ 数据。

FSDP 的通信闭环

flowchart LR
  FWD["**前向**"] --> AG1["All-Gather\n拼出完整参数"] --> COMP1["计算"] --> FREE["释放完整参数"]
  BWD["**反向**"] --> AG2["All-Gather\n拼出完整参数"] --> COMP2["计算梯度"] --> RS["Reduce-Scatter\n只保留梯度分片"]
  UPD["**更新**"] --> LOCAL["每卡用自己的\n梯度分片更新\n自己的参数分片"]

关键洞察：All-Reduce 在概念上等价于 Reduce-Scatter + All-Gather。DDP 用 All-Reduce 是因为每张卡需要完整梯度；FSDP 用 Reduce-Scatter 是因为每张卡只需要自己那份。NCCL 在实现 All-Reduce 时，内部也经常将其分解为这两步。

场景三：Pipeline Parallel — 模型按层切开，激活值怎么传？

Pipeline Parallel（PP） 把模型按层划分为多个 stage，每个 stage 放在不同的 GPU 上。前向传播时，stage 0 算完要把激活值传给 stage 1；反向传播时，stage 1 要把梯度传回 stage 0。

这不需要集合通信——就是两张卡之间直接传数据 → Send / Recv。

Send / Recv（Point-to-Point）

功能：两张卡之间的点对点通信。

flowchart LR
  R0["Rank 0"] -->|"send(activations) — 前向"| R1["Rank 1"]
  R1 -->|"send(gradients) — 反向"| R0

通信量：$O(N)$，仅涉及两张卡。

Send/Recv 是唯一的非集合通信原语。Pipeline 调度算法（1F1B、Zero Bubble 等）的本质就是精心编排这些 Send/Recv 的时序，让不同 stage 尽量同时忙碌，减少流水线气泡。

场景四：MoE Expert Parallel — token 路由到不同专家，怎么搬？

在 MoE（Mixture-of-Experts） 模型中，每个 token 经过 gating network 被路由到一个或几个专家。当使用 Expert Parallel（EP） 时，不同的专家分布在不同的 GPU 上。

问题是：每张卡上的 token 可能需要去任意一张卡上的专家。这不是"一对多"或"多对一"，而是每张卡都要向每张卡发送不同的数据 → All-to-All。

All-to-All

功能：每张卡向其他所有卡发送不同的数据块，同时接收来自所有卡的数据块。

flowchart LR
  subgraph Before["Before — 按行: 每卡的数据"]
    direction TB
    R0B["Rank 0: a→0 a→1 a→2 a→3"]
    R1B["Rank 1: b→0 b→1 b→2 b→3"]
    R2B["Rank 2: c→0 c→1 c→2 c→3"]
    R3B["Rank 3: d→0 d→1 d→2 d→3"]
  end
  Before -->|"all_to_all"| After
  subgraph After["After — 按列: 每卡收集来自所有卡的数据"]
    direction TB
    R0A["Rank 0: a→0 b→0 c→0 d→0"]
    R1A["Rank 1: a→1 b→1 c→1 d→1"]
    R2A["Rank 2: a→2 b→2 c→2 d→2"]
    R3A["Rank 3: a→3 b→3 c→3 d→3"]
  end

MoE 层的通信模式是：All-to-All（dispatch: token → expert）→ 专家计算 → All-to-All（combine: expert output → 原始卡），一前一后两次 All-to-All。

通信量：每张卡发送和接收 $N \cdot \frac{P-1}{P}$ 数据。All-to-All 是最"重"的集合通信，因为它涉及全网状（full-mesh）的数据交换，对网络拓扑和带宽极为敏感。

场景五：初始化与数据流 — 一些"胶水"原语

上面四个场景覆盖了训练中的核心通信需求。还有几个原语用于初始化和数据管理：

Broadcast

功能：把一张卡上的数据广播到所有卡。

flowchart LR
  subgraph Before[" "]
    direction TB
    R0B["Rank 0: A A A A"]
    R1B["Rank 1: . . . ."]
    R2B["Rank 2: . . . ."]
    R3B["Rank 3: . . . ."]
  end
  Before -->|"broadcast\n(from Rank 0)"| After
  subgraph After[" "]
    direction TB
    R0A["Rank 0: A A A A"]
    R1A["Rank 1: A A A A"]
    R2A["Rank 2: A A A A"]
    R3A["Rank 3: A A A A"]
  end

典型场景：训练开始前，Rank 0 上初始化模型参数，通过 Broadcast 确保所有卡参数一致。DDP 启动时内部就会调用 Broadcast。

通信量：每张卡发送或接收 $O(N)$ 数据。但全局总通信量取决于实现：朴素实现（Root 逐一发送）为 $O(N \cdot P)$；实际的树形广播中，多张卡并行转发，总步数为 $\log P$，全局通信量为 $O(N \log P)$。

Scatter / Gather

Scatter：从一张卡分发不同数据块到各卡。Gather：各卡的数据收集到一张卡（Scatter 的逆操作）。

flowchart LR
  subgraph SCATTER["Scatter (Rank 0 分发)"]
    direction LR
    S_IN["Rank 0: d0 d1 d2 d3"] -->|"scatter"| S_OUT0["Rank 0: d0"]
    S_IN --> S_OUT1["Rank 1: d1"]
    S_IN --> S_OUT2["Rank 2: d2"]
    S_IN --> S_OUT3["Rank 3: d3"]
  end
  subgraph GATHER["Gather (收集到 Rank 0)"]
    direction LR
    G_IN0["Rank 0: d0"] -->|"gather"| G_OUT["Rank 0: d0 d1 d2 d3"]
    G_IN1["Rank 1: d1"] --> G_OUT
    G_IN2["Rank 2: d2"] --> G_OUT
    G_IN3["Rank 3: d3"] --> G_OUT
  end

典型场景：数据加载时，Rank 0 读取一个大 batch 然后 Scatter 到各卡；评估阶段 Gather 各卡的预测结果到 Rank 0 做汇总。

全景回顾：从场景到原语

以下代码展示了如何在 PyTorch 中使用 all_reduce：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

import torch
import torch.distributed as dist
import os

def main():
    dist.init_process_group(backend="nccl")
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    device = torch.device(f"cuda:{rank}")

    # 每张卡创建自己的 tensor
    tensor = torch.ones(1024, 1024, device=device) * (rank + 1)
    print(f"[Rank {rank}] Before all_reduce: sum = {tensor.sum().item():.0f}")

    # All-Reduce: 求和
    dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
    print(f"[Rank {rank}] After all_reduce: sum = {tensor.sum().item():.0f}")
    # 期望结果: 每个元素 = 1+2+3+4 = 10 (4 卡)
    # 总和 = 10 * 1024 * 1024 = 10,485,760

    dist.destroy_process_group()

if __name__ == "__main__":
    main()

# 运行: torchrun --nproc_per_node=4 nccl_allreduce.py

完整代码（包含所有原语的示例）见 code/01-gpu-memory-distributed/nccl_allreduce.py

集合通信算法

通信原语定义了"做什么"，通信算法决定了"怎么做"。同一个 all_reduce 操作，不同的算法在延迟和带宽利用上差异巨大。

Ring All-Reduce

Ring All-Reduce 是最经典的带宽最优算法，分为两个阶段：

阶段一：Reduce-Scatter（$P-1$ 步）

4 张卡排成一个环，每张卡将数据分为 $P=4$ 个块。每一步，每张卡向下一个邻居发送一个块，同时从上一个邻居接收一个块并累加。经过 $P-1=3$ 步后，每张卡上有一个块包含了所有卡的归约结果。

flowchart TD
  subgraph INIT["初始状态"]
    direction LR
    I0["Rank 0: a0 a1 a2 a3"] ~~~ I1["Rank 1: b0 b1 b2 b3"]
    I2["Rank 2: c0 c1 c2 c3"] ~~~ I3["Rank 3: d0 d1 d2 d3"]
  end

  subgraph STEP1["Step 1 — 每卡发送一个块给右邻居，接收左邻居的块并累加"]
    direction LR
    S1R0["Rank 0: a0 · a1 · a2 · **a3+d3**"]
    S1R1["Rank 1: **b0+a0** · b1 · b2 · b3"]
    S1R2["Rank 2: c0 · **c1+b1** · c2 · c3"]
    S1R3["Rank 3: d0 · d1 · **d2+c2** · d3"]
  end

  subgraph STEP3["Step 3 (最终) — 每张卡上有一个完整归约的块"]
    direction LR
    S3R0["Rank 0: · · · · · · **Σ3** ← 块 3 完整"]
    S3R1["Rank 1: **Σ0** · · · · · · ← 块 0 完整"]
    S3R2["Rank 2: · **Σ1** · · · · ← 块 1 完整"]
    S3R3["Rank 3: · · **Σ2** · · ← 块 2 完整"]
  end

  INIT --> STEP1 -->|"继续传递\n部分归约的块"| STEP3

阶段二：All-Gather（$P-1$ 步）

同样在环上传递，但这次不做归约，只做拷贝。$P-1$ 步后每张卡拥有完整的归约结果。

复杂度分析：

• 总步数：$2(P-1)$
• 每步每卡传输量：$N/P$
• 总通信量（每卡）：$2 \cdot \frac{P-1}{P} \cdot N$
• 当 $P$ 很大时趋近于 $2N$，与 GPU 数量无关——这就是带宽最优的含义

Ring 的缺点是延迟为 $O(P)$：数据必须绕环一圈，每一步都有一次网络延迟。当消息较小时，延迟开销会超过传输时间，效率降低。

Tree All-Reduce

Tree All-Reduce 用一棵二叉树组织通信：

graph TD
  R0["Rank 0 (Root)"]
  R1["Rank 1"]
  R2["Rank 2"]
  R3["Rank 3"]

  R0 --- R1
  R0 --- R2
  R1 --- R3

  style R0 fill:#d4a574,color:#000

阶段一 (Reduce): 叶子向根归约 | 阶段二 (Broadcast): 根向叶子广播

复杂度：

• 延迟：$O(\log P)$——远优于 Ring 的 $O(P)$
• 带宽利用率：较差，非叶子节点成为带宽瓶颈

Tree 适合小消息、多节点的场景。实际中 NCCL 会根据消息大小自动选择算法。

Recursive Halving-Doubling

这是一种折中方案，结合了 Ring 的带宽效率和 Tree 的低延迟：

1. Halving 阶段：每轮将参与者分成两半，两半之间交换并归约各自缺少的部分
2. Doubling 阶段：反向传播完整结果

延迟为 $O(\log P)$，带宽利用率接近最优。适合 GPU 数为 2 的幂次的情况。

NCCL 的实际选择策略：NCCL 并不固定使用某一种算法，而是根据消息大小、GPU 数量、拓扑结构动态选择：

• 小消息（< 256 KB）：偏向 Tree
• 大消息（> 数 MB）：偏向 Ring
• 特定拓扑下会使用更高效的变种

通信拓扑

通信算法跑在硬件拓扑之上。不同的硬件链路带宽差异巨大，直接影响了分布式训练的瓶颈位置。

NVLink

NVLink 是 NVIDIA GPU 之间的高速直连链路：

A100 机内 8 卡通过 NVLink 互联，每对 GPU 间带宽 600 GB/s——是 HBM 带宽的 30%，是 PCIe Gen4 x16 的 20 倍以上。

NVSwitch

NVSwitch 是 NVIDIA 的全交换芯片，实现了节点内 GPU 之间的全双工（Full Bisection Bandwidth） 互联：

flowchart LR
  subgraph DGX["DGX A100 — 任意 GPU 对: 600 GB/s, 聚合: 4.8 TB/s"]
    direction LR
    subgraph LEFT[" "]
      direction TB
      G0["GPU 0"] ~~~ G1["GPU 1"] ~~~ G2["GPU 2"] ~~~ G3["GPU 3"]
    end
    NVS["NVSwitch\n× 6"]
    subgraph RIGHT[" "]
      direction TB
      G4["GPU 4"] ~~~ G5["GPU 5"] ~~~ G6["GPU 6"] ~~~ G7["GPU 7"]
    end
    LEFT <-->|"NVLink"| NVS <-->|"NVLink"| RIGHT
  end

有了 NVSwitch，节点内 all_reduce 的带宽几乎不受 GPU 对数限制。

PCIe

PCIe 是 CPU 和 GPU 之间、以及没有 NVLink 的 GPU 之间的通信通道：

• PCIe Gen4 x16: 约 32 GB/s（双向）
• PCIe Gen5 x16: 约 64 GB/s（双向）

相比 NVLink 600 GB/s，PCIe 带宽低了一个数量级。在消费级 GPU 或部分云实例上，GPU 间通信可能退化到走 PCIe，此时通信会成为严重瓶颈。

RDMA / InfiniBand

跨节点通信依赖网络互联，主流方案是 InfiniBand + GPUDirect RDMA：

• InfiniBand HDR: 200 Gb/s = 25 GB/s
• InfiniBand NDR: 400 Gb/s = 50 GB/s
• GPUDirect RDMA: GPU 显存直接通过网卡发送数据，绕过 CPU 和系统内存，减少一次拷贝延迟

flowchart LR
  subgraph A["节点 A — NVLink 600 GB/s"]
    direction TB
    A1["GPU"] <-->|NVLink| A2["GPU"]
    A3["GPU"] <-->|NVLink| A4["GPU"]
    A1 <-->|NVSwitch| A3
    NICA["NIC"]
  end
  subgraph B["节点 B — NVLink 600 GB/s"]
    direction TB
    B1["GPU"] <-->|NVLink| B2["GPU"]
    B3["GPU"] <-->|NVLink| B4["GPU"]
    B1 <-->|NVSwitch| B3
    NICB["NIC"]
  end
  NICA <====>|"InfiniBand RDMA\n25-50 GB/s"| NICB

节点内外的带宽差距（约 10-20 倍）深刻影响了并行策略的设计：

• 通信量大的并行策略（如 TP）通常放在节点内，充分利用 NVLink
• 通信量较小的并行策略（如 PP、DP）可以跨节点部署
• 这就是大规模训练中"节点内 TP、节点间 DP/PP"成为标准配置的原因

总结与下一步

本文覆盖了 LLM Infra 的两大基石：

显存方面：

• GPU 采用吞吐优先的设计，通过大量简单核心和 Warp 级别的延迟隐藏实现高吞吐
• 内存层级从寄存器到 HBM，带宽跨越 4 个数量级
• LLM 推理（自回归解码）是典型的 Memory-bound 问题，算术强度仅约 1 FLOPs/Byte
• 训练时，Adam 优化器让显存需求膨胀到参数量的 16 倍

通信方面：

• 8 种 NCCL 通信原语各有其对应的分布式并行场景
• Ring All-Reduce 带宽最优但延迟 $O(P)$，Tree 延迟最优但带宽差
• 节点内 NVLink (600 GB/s) 与节点间 InfiniBand (25-50 GB/s) 的带宽鸿沟，决定了混合并行的拓扑布局

在下一篇文章**《分布式并行策略全景》**中，我们将在这些基础之上，系统介绍 DDP、FSDP、TP、PP、SP、EP 等并行策略——你会看到，每种策略本质上都是在显存和通信之间做不同的 trade-off，而选择哪种 trade-off，取决于模型规模、硬件拓扑和训练/推理场景。

参考资料

1. NVIDIA CUDA Programming Guide — docs.nvidia.com/cuda — SM、Warp、内存层级的权威文档
2. NVIDIA A100 Whitepaper — GPU 架构细节、Tensor Core 规格、NVLink/NVSwitch 拓扑
3. NCCL Documentation — docs.nvidia.com/deeplearning/nccl — 通信原语 API 和算法说明
4. Roofline Model — Williams, Waterman, Patterson, “Roofline: An Insightful Visual Performance Model for Multicore Architectures”, Communications of the ACM, 2009
5. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models — Rajbhandari et al., 2020 — 对优化器状态、梯度、参数的显存分析
6. Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM — Narayanan et al., 2021 — 混合并行策略与通信拓扑的关系
7. PyTorch Distributed Overview — pytorch.org/tutorials — PyTorch 分布式通信接口