CS336 Lecture Notes 4

集体操作（Collective Operations）#

集体操作是分布式编程的概念原语，源自1980年代的并行编程文献。“集体”意味着指定跨多个设备的一般通信模式。

基础概念：

• Rank：特定设备/GPU的标识（如0, 1, 2, 3）
• World size：设备总数（如4）

基础操作： Broadcast，Scatter，Gather，Reduce

核心操作：：AllGather，ReduceScatter，AllReduce，AlltoAll

硬件基础#

家庭环境：

• 同节点GPU通过PCIe总线通信（v7.0, 16通道 → 242 GB/s）[wiki]
• 不同节点通过以太网通信（~200 MB/s）

现代数据中心：

典型配置：

• 每节点 8 GPU，通过 NVLink 连接到 NVSwitch（B200 的 NVLink 5.0 达 1.8 TB/s；HBM 为8 TB/s）
• 每 pod 256 节点，通过 Infiniband 连接（~0.05 TB/s）
• N 个 pod 组成集群/数据中心，通过以太网连接（通过 PCIe -> CPU）

绕过CPU：

• 以太网需要经过 CPU（复制数据到内核 socket 缓冲区，构建 TCP 包，复制到 NIC 环形缓冲区）
• RDMA（远程直接内存访问）：允许一个 GPU 直接读写另一个 GPU 的内存，无需 CPU 参与
• Infiniband 支持 RDMA，但标准以太网不支持

新发展：

• GB200/GB300 NVL72：每 tray 8 GPU，每 rack 9 tray -> 一个 NVLink 域内 72 GPU
• RoCE（RDMA over Converged Ethernet）：以太网绕过 CPU，类似 Infiniband 但更便宜/弱，Meta 使用 RoCE

NCCL（NVIDIA Collective Communication Library）#

NCCL将集体操作翻译为 GPU 间发送的低级数据包：[talk]

• 检测硬件拓扑（节点数、交换机、NVLink/PCIe）
• 优化 GPU 间路径
• 启动 GPU Kernels 发送/接收数据

PyTorch Distributed#

PyTorch 分布式库（torch.distributed）[documentation]

• 为集合通信操作（例如 all_gather_into_tensor）提供简洁易用的接口
• 支持适配不同硬件的多种后端：gloo（适配 CPU）、nccl（适配 GPU）
• 同时支持高阶分布式算法（例如 FullyShardedDataParallel）

代码详见 [集合通信操作]。

带宽基准测试#

参考：

• [How to reason about collective operations]
• [Sample benchmarking code]

代码详见 [带宽基准测试]

All-reduce有效带宽计算：

1
sent_bytes = size_bytes * 2 * (world_size - 1)  # 2x 因为发送 + 接收
2
total_duration = world_size * duration
3
bandwidth = sent_bytes / total_duration

关键发现：All-reduce = Reduce-scatter + All-gather，传输 2 倍数据花费 2 倍时间，带宽相近。

数据并行（Data Parallelism）#

分片策略： 每个 rank 获取数据的一个切片。

随机生成一个 shape 为 (batch_size, num_dim) 的张量作为数据输入。

1
def generate_sample_data():
2
    batch_size = 128
3
    num_dim = 1024
4
    data = torch.randn(batch_size, num_dim)
5
    return data
6

7
data = generate_sample_data()

每个 rank 获取自己的一份输入。

1
# 原始数据批次 → 分片到各GPU
2
# --- B0 ---   →  rank 0
3
# --- B1 ---   →  rank 1
4
# --- B2 ---   →  rank 2
5
# --- B3 ---   →  rank 3
6
local_batch_size = int_divide(batch_size, world_size)
7
start_index = rank * local_batch_size
8
end_index = start_index + local_batch_size
9
data = data[start_index:end_index].to(cuda_if_available(rank))

1
def int_divide(a: int, b: int):
2
    """Return a / b and throw an error if there's a remainder."""
3
    assert a % b == 0
4
    return a // b

参数获取，优化器设置以及训练循环：

1
# 创建 MLP parameters, params[0], ..., params[num_layers - 1], 每个 rank 持有完整参数
2
params = [get_init_params(num_dim, num_dim, rank) for layer in range(num_layers)]
3
optimizer = torch.optim.AdamW(params, lr=1e-3)  # 每个 rank 都有自己的 optimizer state
4

5
for step in range(num_steps):
6
    # 前向传播（各rank独立计算）
7
    x = data
8
    for param in params:
9
        x = x @ param
10
        x = F.gelu(x)
11
    loss = x.square().mean()
12

13
    # 反向传播
14
    loss.backward()
15

16
    # 关键: 同步梯度（DDP与标准训练的唯一区别）
17
    for param in params:
18
        dist.all_reduce(tensor=param.grad, op=dist.ReduceOp.AVG)
19

20
    # 更新参数
21
    optimizer.step()
22

23
    print(f"[data_parallelism] Rank {rank}: step = {step}, loss = {loss.item()}, params = {[summarize_tensor(params[layer]) for layer in range(num_layers)]}", flush=True)

1
def get_init_params(num_inputs: int, num_outputs: int, rank: int) -> nn.Parameter:
2
    """Create parameters and put them on the `rank`-th GPU."""
3
    torch.random.manual_seed(0)  # For reproducibility
4
    return nn.Parameter(torch.randn(num_inputs, num_outputs, device=cuda_if_available(rank)) / math.sqrt(num_outputs))

1
def summarize_tensor(tensor: tensor) -> str:
2
    return "x".join(map(str, tensor.shape)) + "[" + str(round(tensor.view(-1)[0].item(), 4)) + "...]"

特点：

• 各 rank 的 loss 不同（基于本地数据计算）
• 梯度通过 all-reduce 保持一致
• 参数在各 rank 保持相同

局限性： 每个 rank 需持有全部参数，内存压力大

张量并行（Tensor Parallelism）#

分片策略： 每个 rank 获取每层的部分参数

每个 rank 都收到全部的输入 data（batch_size x num_dim）

1
data = data.to(cuda_if_available(rank))  # All ranks get the data (batch_size x num_dim)

不包含 backward 的训练循环：

1
# 每层权重矩阵 W 被按列切分：
2
#  |  |  |  |
3
#  W0 W1 W2 W3  （各 rank 持有 1/world_size 的参数）
4
#  |  |  |  |
5
local_num_dim = int_divide(num_dim, world_size)  # Shard `num_dim`
6
params = [get_init_params(num_dim, local_num_dim, rank) for layer in range(num_layers)]
7

8
# 前向传播
9
x = data  # 所有 rank 获得完整数据
10
for layer in range(num_layers):
11
    x = x @ params[layer]  # 仅在参数切片上计算
12
    x = F.gelu(x)
13

14
    # 为激活分配内存 (world_size x batch_size x local_num_dim)
15
    activations = [torch.empty(batch_size, local_num_dim, device=cuda_if_available(rank)) for _ in range(world_size)]
16

17
    # All-gather 收集所有 rank 的激活
18
    dist.all_gather(tensor_list=activations, tensor=x)
19

20
    # 拼接得到完整激活 batch_size x num_dim
21
    x = torch.cat(activations, dim=1)
22

23
print(f"[tensor_parallelism] Rank {rank}: forward pass produced activations {summarize_tensor(x)}", flush=True)
24

25
# backward ... ?

特点：

• 每层需要 all-gather 传递激活
• 需要非常快的互连（如 NVLink）
• 适合大模型单节点内并行

流水线并行（Pipeline Parallelism）#

分片策略： 每个 rank 获取部分层

1
data = data.to(cuda_if_available(rank))
2

3
# 按层进行切分
4
local_num_layers = int_divide(num_layers, world_size)
5

6
# 每个 rank 获取一部分层
7
local_params = [get_init_params(num_dim, num_dim, rank) for layer in range(local_num_layers)]
8

9
# 前向传播
10

11
# 将批次切分为微批次以减少气泡（bubble）
12
micro_batch_size = int_divide(batch_size, num_micro_batches)
13
if rank == 0:
14
    # 训练数据
15
    micro_batches = data.chunk(chunks=num_micro_batches, dim=0)
16
else:
17
    # 为激活值分配内存
18
    micro_batches = [torch.empty(micro_batch_size, num_dim, device=cuda_if_available(rank)) for _ in range(num_micro_batches)]
19

20
for x in micro_batches:
21
    # 从前一 rank 接收激活
22
    if rank - 1 >= 0:
23
        dist.recv(tensor=x, src=rank - 1)
24

25
    # 计算分配给本 rank 的层
26
    for param in local_params:
27
        x = x @ param
28
        x = F.gelu(x)
29

30
    # 发送到下一 rank
31
    if rank + 1 < world_size:
32
        print(f"[pipeline_parallelism] Rank {rank}: sending {summarize_tensor(x)} to rank {rank + 1}", flush=True)
33
        dist.send(tensor=x, dst=rank + 1)
34

35
# backward ... ?

特点：

• 可在慢速互连下工作
• 需要处理流水线气泡（通过微批次）
• 点对点通信（send/recv）

TPU vs GPU 网络设计差异#

为什么不全连接？

• Domain sizes限制：全连接成本过高
• TPU8i/t 转向树型拓扑（可能为MoE优化）
• TPU8t 使用交换网络（‘Virgo’）

朴素数据并行（Naïve Data Parallelism）#

核心思想： 将B大小的批次切分到M台机器，交换梯度同步

1
𝜃_{t+1} = 𝜃_t - 𝜂 * Σ ∇f(x_i)  # SGD公式

性能分析：

• 计算扩展：每个GPU处理 B/M 样本
• 通信开销：每批次传输 2×#params
• 内存扩展：无！每个GPU需要完整参数

内存问题深度分析#

内存实际需求（每参数）：

需要5份权重副本，内存开销巨大！

ZeRO：解决DP内存开销#

核心思想： 分片昂贵的部分（状态），利用reduce-scatter等价性

ZeRO Stage 1（优化器状态分片）：

1
工作流程：
2
Step 1: 各GPU计算完整梯度
3
Step 2: ReduceScatter梯度 → 通信量 #params
4
Step 3: 各机器更新其负责的参数切片
5
Step 4: AllGather参数 → 通信量 #params

结论：ZeRO Stage 1 是免费的内存优化！

ZeRO Stage 2（梯度分片）：

进一步分片梯度：

• 梯度也跨机器分片
• 使用增量计算：计算一层梯度后立即reduce发送到正确worker
• 不再需要完整的梯度向量

ZeRO Stage 3 / FSDP（参数也分片）：

1
工作流程（简化版）：
2
- 分片所有内容（包括参数）
3
- 按需请求/发送参数
4
- 通信成本：2次 all-gather + 1次 reduce-scatter = 3×#params

关键技巧：

1. 增量计算/通信：参数/梯度请求后立即释放
2. 通信计算重叠：all-gather与前向传播同时进行

ZeRO 实践：模型容量估算#

纯BF16训练（含Kahan累加）：12 bytes/param

8×A100 80GB 的最大模型容量：

数据并行局限性#

计算扩展问题：

• #machines < batch_size（接近此限制时通信开销高）
• 大批次收益递减

内存扩展问题：

• ZeRO-1/2 不能扩展内存
• ZeRO-3 不减少激活内存

需要更好的模型分片方案……

流水线并行（Pipeline Parallelism）#

问题：层级并行利用率极低

1
简单层级并行：
2
- n个GPU，每个GPU仅 1/n 时间活跃
3
- 大部分时间空闲，等待反向传播

解决方案：微批次流水线

1
处理4个微批次：
2
- 发送第一个微批次后立即计算第二个
3
- 气泡时间比例 = (n_stages - 1) / n_micro
4
- 需要大批次来隐藏气泡！

为什么使用流水线并行？

1. 节省内存（相比DDP）
2. 通信特性好（相比FSDP）—— 仅激活 (b×s×h)，点对点

适用场景： 慢速网络链接（跨节点）

气泡优化策略：

张量并行（Tensor Parallelism）#

核心观察：矩阵乘法可分解

1
矩阵乘法分解：
2
C = AB = (A1 + A2)(B1 + B2) = A1B1 + A1B2 + A2B1 + A2B2
3

4
列并行 + 行并行：
5
- 前向传播：f=identity, g=all-reduce
6
- 反向传播：f=all-reduce, g=identity

Transformer中的张量并行：

何时使用张量并行？

• GPU节点内（最多8 GPU）—— 高速互连

张量并行 vs 流水线并行：

激活内存问题#

激活内存每层需求（存储所有）：

1
激活内存公式：
2
- 二次注意力项（含dropout）：5as/h
3
- 可通过Flash Attention重计算消除
4

5
剩余项（10sbh）：
6
- LayerNorm: 4sbh
7
- Dropout: 2sbh
8
- Attention/MLP输入: 4sbh

序列并行（Sequence Parallelism）#

核心观察：10sbh项都是序列上的逐点操作

1
解决方案：
2
- 沿序列轴切分LayerNorm/Dropout
3
- 前向：'g' = all-gather, 'ḡ' = reduce-scatter
4
- 反向：两者互换

效果： 激活内存也随机器数线性扩展！

专家并行（Expert Parallelism）#

核心思想： 不切分矩阵乘法，切分专家并路由激活

1
MoE结构：
2
- MLP使用专家并行
3
- Attention使用张量并行
4

5
复杂问题：
6
- EP和DP共享副本（EP < DP）
7
- DP和TP可能交互降低利用率

Attention vs MLP并行解耦：

Megatron灵活策略：

• Attention: TP/CP/DP
• MLP: ETP/EP/EDP

其他并行策略#

Context Parallel / Ring Attention：

• 长序列场景下，将激活跨GPU切分

3D/4D并行组合规则#

简单经验法则：

1
1. 直到模型适配内存：
2
   - Tensor/Expert并行 → GPU数/机器（最多8）
3
   - Pipeline并行 → 跨机器
4
   - （或用ZeRO-3，取决于带宽）
5

6
2. 直到GPU耗尽：
7
   - 其余用数据并行扩展
8

9
如果批次小：梯度累加换取更大批次

Megatron推荐配置示例：

1
基础：TP + PP + DP
2
高级：TP + PP + CP + EP（MoE）

Narayanan 2021扩展策略#

1
配置趋势：
2
- Tensor并行先扩展到8，然后封顶在8
3
- Pipeline并行扩展直到模型适配
4
- 数据并行随规模逐渐减少，最大模型DP=6

关键发现：

• 64机器配置最优：8×8（TP=8）
• 精心设计的3D并行可获得线性增益

激活重计算的收益：

• 激活重计算 → 更大批次 → 更高吞吐

实际LLM训练案例#

模式总结：

• TP通常 ≤ 8
• EP可以很大（但实现难！）
• 长上下文阶段使用大CP

Llama 3 405B 训练细节#

1
阶段配置：
2
- Stage 1：小批次训练
3
- Stage 2：预训练
4
- Stage 3：长上下文
5

6
挑战：此规模下GPU故障频繁！