1.1 数字对比

GPU之间通信，走什么路？两条路：

NVLink 4.0的总带宽是PCIe 5.0的7倍以上。在DGX H100的8卡全互联拓扑中，总带宽可达7.2 TB/s。到了NVLink 5.0（B300），这个数字飙升至1.8 TB/s。

1.2 为什么差距这么大？

PCIe是通用总线，设计目标是连接CPU与各种外设（GPU、网卡、SSD等）。数据从GPU A到GPU B的路径是：

1
GPU A → PCIe Switch → CPU内存（DMA）→ PCIe Switch → GPU B

每一步都有开销，CPU还要参与地址转换和中断处理。

NVLink是GPU专用直连总线，设计目标只有一个：让GPU之间直接通信，不需要CPU当中间人。NVLink支持统一内存访问——GPU可以直接读写对方显存，无需CPU介入。

1.3 实践意义

• NVLink适合节点内（同一台机器内的8张卡）：张量并行（TP）强烈依赖高频通信，应放在NVLink域内
• PCIe适合节点间（跨机器）：数据并行（DP）每个step只同步一次梯度，对带宽不那么敏感

3.1 核心通信原语

All-Reduce：分布式训练中最核心的原语。将N个GPU上的数据进行归约（求和、取最大等），然后把结果广播到所有GPU。

数学上，对于K个rank，每个rank提供数组 $V_k$，All-Reduce后每个rank都得到：

$S[i] = V_0[i] + V_1[i] + \cdots + V_{K-1}[i]$

All-Gather：从所有rank收集数据，每个rank得到完整的 $K \times N$ 结果。

Reduce-Scatter：归约后将结果分散到各个rank，每个rank只拿到自己那份。

关键关系：Reduce-Scatter + All-Gather = All-Reduce。

3.2 Ring All-Reduce：带宽最优

NCCL根据数据量自动选择算法：

• 小数据（<32KB） ：Tree算法，$O(\log_2 N)$ 步，延迟最优
• 中等数据（32KB-2MB） ：Double Binary Tree，平衡延迟和带宽
• 大数据（>2MB） ：Ring算法，带宽最高

Ring All-Reduce的原理：

1. Reduce-Scatter阶段：N个GPU构成逻辑环，沿环传递N-1轮，每轮每个GPU发送 $1/N$ 的数据并归约接收到的数据。完成后每个GPU持有 $1/N$ 的完整归约结果。

2. All-Gather阶段：再沿环传递N-1轮，每轮发送自己的归约结果。完成后所有GPU拥有完整结果。

通信量：每个GPU单向发送 $2 \times (N-1)/N \times \text{data\_size}$。8 GPU时约为 $1.75 \times \text{data\_size}$。带宽效率高，但延迟随N线性增长。

Tree All-Reduce只需 $O(\log_2 N)$ 步——8 GPU时3步完成，远快于Ring的14步。但带宽效率不如Ring。

3.3 NCCL代码示例

1
import torch
2
import torch.distributed as dist
3

4
# 初始化进程组
5
dist.init_process_group(backend='nccl')
6

7
# 假设每个rank有张量tensor
8
tensor = torch.randn(1024, 1024).cuda()
9

10
# All-Reduce：所有GPU求和，结果广播到所有GPU
11
dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
12

13
# All-Gather：收集所有GPU的数据
14
gather_list = [torch.zeros_like(tensor) for _ in range(dist.get_world_size())]
15
dist.all_gather(gather_list, tensor)
16

17
# Reduce-Scatter：归约后分散
18
dist.reduce_scatter(output_tensor, input_list, op=dist.ReduceOp.SUM)

4.1 数据并行（Data Parallelism, DP）

原理：每张GPU存一份完整模型，把数据切成多份，各卡处理自己的数据分片，梯度通过All-Reduce同步。

通信量：每个训练step一次All-Reduce，数据量为模型参数量的2倍（梯度大小=参数量）。

优点：实现简单，通信量不随模型深度增加。

缺点：每张卡都要存完整模型，无法突破单卡显存限制。

适用场景：模型能放进单卡，想提升吞吐量。

4.2 模型并行 / 张量并行（Model Parallelism / Tensor Parallelism, MP/TP）

原理：把每一层的权重矩阵切分到多张GPU上。以Transformer的Attention层为例，将 $[\text{hidden\_dim}, \text{num\_heads} \times \text{head\_dim}]$ 的权重矩阵沿列方向切开。

通信量：非常高——每层前向和反向都需要All-Reduce同步。

优点：突破单卡显存限制，适合超大模型。

缺点：通信频繁（每层都同步），通信占比可达30%+。

适用场景：节点内（NVLink高速互联），单卡放不下单层。

4.3 流水线并行（Pipeline Parallelism, PP）

原理：把模型的层切分到不同GPU——GPU 0负责前N层，GPU 1负责接下来的N层，依此类推。前向传播时数据从左向右流经各卡。

通信量：相对较小——只在相邻设备间传递激活值（前向）和梯度（反向），不是完整的梯度同步。

优点：通信量小，可跨节点扩展。

缺点：存在流水线气泡（bubble）——部分GPU空闲等待。

适用场景：节点间，模型极深（千亿参数以上）。

4.4 通信量对比总结

最佳实践：节点内做TP，节点间做DP——TP走NVLink高速通道，DP走PCIe/网络。

5.1 模型状态的构成

以Adam优化器为例，训练时每张GPU需要存储：

• 参数（Parameters） ：$P$
• 梯度（Gradients） ：$P$
• 优化器状态（Optimizer States） ：$2P$（Adam的动量和方差）

总计：$4P$

5.2 ZeRO三个阶段

ZeRO-1：将优化器状态（Adam动量）分片到N个GPU。单卡显存从 $4P$ 降至 $P + P + 2P/N \approx (2 + 2/N)P$。8卡时减少约62.5%。

ZeRO-2：进一步分片梯度。单卡显存降至 $P + P/N + 2P/N = (1 + 3/N)P$。8卡时减少约75%。

ZeRO-3：分片一切——参数、梯度、优化器状态全部打散。单卡显存降至 $4P/N$。8卡时减少约87.5%。

代价：每增加一个stage，通信开销就增加一层。ZeRO-3需要在前向传播前All-Gather参数，在反向传播后Reduce-Scatter梯度。

5.3 FSDP：PyTorch的ZeRO-3实现

FSDP（Fully Sharded Data Parallel） 是PyTorch对ZeRO-3的原生实现。它将所有模型状态（参数、梯度、优化器状态）分片到所有数据并行worker上，达到理论上的最小单卡显存。

FSDP的训练流程：

1. 前向传播前：All-Gather当前层的参数分片 → 所有GPU获得完整参数
2. 前向计算：用完整参数计算
3. 前向传播后：释放非本地的参数分片（reshard_after_forward=True）
4. 反向传播：类似地All-Gather参数，计算梯度
5. 梯度同步：Reduce-Scatter梯度，每张卡只保留自己那部分

FSDP支持ZeRO-2和ZeRO-3两种模式。