Runguo Li

Runguo Li

M.S. student in Information Science @ UIUC, advised by Prof. Minjia Zhang. Working on ML systems for large-scale models: MoE inference and expert streaming, distributed training and rollout, GPU compiler correctness.

第 7 讲：多 GPU 并行

目录 · ← l6 · l8 →

第 7 讲：多 GPU 并行

日期：4 月 20 日（周一，Spring 2026） | 讲师：Percy Liang | 材料：lecture_07.py

概览

上周讲的是单卡内部的并行（融合、分块），本周讲跨 GPU 的并行：集合通信操作（broadcast、scatter、gather、reduce、all-gather、reduce-scatter、all-reduce、all-to-all）、硬件连接（NVLink/NVSwitch vs Infiniband vs Ethernet、RDMA），以及用 torch.distributed 在深层 MLP 上从零实现三种经典策略——数据并行（切 batch）、张量并行（切宽度）、流水线并行（切深度）。

核心概念与定义

  • 统一主题：计算（ALU）离数据太远。单卡内靠融合/分块减少 HBM 访问；跨卡靠复制/分片减少网络流量。层次结构：L1/共享内存（最快）→ HBM → NVLink/NVSwitch（节点内）→ Infiniband/Ethernet（节点间，最慢）。
  • 为什么要多 GPU？ (1) 参数 + 优化器状态 + 梯度 + 激活放不进单卡；(2) 更多 GPU = 更多 FLOPs = 训练更快。
  • 集合通信操作（1980 年代并行计算文献中的经典概念；描述的是”跨设备的通信模式“，而非点对点消息）：
    • Rank = 设备编号（0..world_size−1）；world size = 设备总数。
    • Broadcast：从 rank 0 复制到所有 rank（例：rank 0 加载 checkpoint 后广播）。
    • Scatter：把 rank 0 上的一个张量分给各 rank（理解 reduce-scatter 的踏脚石）。
    • Gather：把各 rank 的片段收集到 rank 0（理解 all-gather 的踏脚石）。
    • Reduce：用某个算子（sum/min/max）把各 rank 的数据合并到 rank 0。
    • All-gather：收集到所有 rank（用途：每个 rank 持有参数分片，前向时聚合出完整参数）。
    • Reduce-scatter：先按维度做归约、再把结果散开（用途：反向之后汇总各数据分片的梯度，但把存储分摊出去）。
    • All-reduce = reduce-scatter + all-gather（用途：汇总梯度同时保持完整参数副本——即普通 DDP）。
    • All-to-all：每个 rank 向其他每个 rank 发送一份数据（用途：MoE 的 token 路由；均衡切分时它看起来就是一次转置）。
    • 记忆技巧：reduce = 结合/交换运算；scatter 是 gather 的逆；all = 目的地是所有设备。
  • 硬件：PCIe（家用：242 GB/s）、Ethernet（约 200 MB/s，需经过 CPU）、NVLink→NVSwitch（B200：1.8 TB/s）、Infiniband（约 0.05 TB/s，经 HCA/NIC）。RDMA 让一张 GPU 直接读写另一张 GPU 的显存而不惊动 CPU（Infiniband 支持；标准 Ethernet 不支持；RoCE 是”以太网 + RDMA”，Meta 在用）。NCCL 把集合操作翻译成底层包、探测拓扑并启动收发 kernel。GB200 NVL72 把 72 张 GPU 放在同一个 NVLink 域内。
  • 数据并行（DDP）：每个 rank 拿 batch 的一个切片；每个 rank 持有完整参数副本；本地反向之后 all-reduce(AVG) 梯度，使各 rank 保持同步。
  • 张量并行：每个 rank 持有每层权重的一个切片（例如 MLP 的列切片 W_i）；所有 rank 处理完整 batch；每层结束后 all-gather 部分激活并拼接。因为每层都要通信，需要极快的互联（NVLink）。
  • 流水线并行：每个 rank 持有一部分层（一个 stage）；激活从 rank 0 → 1 → … 流动，用 micro-batch 填满流水线（缩小”气泡”）。能容忍慢互联（点对点、激活大小的通信），但 batch 不大时气泡会拖累性能。
  • 流水线气泡：stage 数为 n、micro-batch 数为 m 时，闲置比例 ≈ (n−1)/m——”所以我们需要很大的 batch！”
  • 本讲未覆盖（明确列出）：通信/计算重叠、attention 相关的并行、序列/专家并行，以及”下次课”的 FSDP/ZeRO（用 all-gather + reduce-scatter 避免持有全部参数）。

代码示例：torch.distributed 中的集合通信

代码（Python）：

import torch.distributed as dist

def collective_operations_main(rank: int, world_size: int):
    setup(rank, world_size)  # init_process_group("nccl"/"gloo", ...)

    data = tensor([0., 1, 2, 3], device=f"cuda:{rank}") + rank
    dist.all_reduce(tensor=data, op=dist.ReduceOp.SUM, async_op=False)  # 原地修改！
    # 之后：每个 rank 都拿到"所有 rank 向量之和"

    input = torch.arange(world_size, dtype=torch.float32) + rank   # 每 rank 一个 [world_size]
    output = torch.empty(1)
    dist.reduce_scatter_tensor(output=output, input=input, op=dist.ReduceOp.SUM)
    # output[rank] = 各 rank 第 rank 列之和

    input = output
    output = torch.empty(world_size)
    dist.all_gather_into_tensor(output_tensor=output, input_tensor=input)
    # output = [第 0 列之和, 第 1 列之和, ...]（每个 rank 都拿到）

    cleanup()  # destroy_process_group()

代码做了什么： 在 world_size 个独立进程（通过 mp.spawn）中运行同一函数；每个进程依次执行 all-reduce、reduce-scatter、all-gather，演示了”all-reduce = reduce-scatter + all-gather”（两者输出一致）。

实现深挖：

  • 为什么是原地操作：dist.all_reduce(tensor=data, ...) 会就地修改 data（既是输入也是输出），避免额外分配——这一点常让人意外。
  • 为什么区分 NCCL 与 gloo：NCCL 是 GPU 后端（会最优利用 NVLink/Infiniband）；gloo 用于 CPU。讲义的 setup 依据 torch.cuda.is_available() 选择。
  • 为什么 async_op=False：同步集合操作会阻塞到完成；要做重叠就得用 async_op=True 并收集 handle——作业 2 明确要求实现异步重叠。
  • 为什么 world_size=4 时用 spawn：每个 rank 是独立 OS 进程、各占一张 GPU；通过 MASTER_ADDR/MASTER_PORT 由 rank 0 协调。

与作业的联系：作业 2 的分布式部分——DDP、优化器状态分片、FSDP——都建立在这些原语之上：梯度 all-reduce（DDP）、reduce-scatter（状态分片）、all-gather + reduce-scatter（FSDP）。讲义测量带宽的代码（all-reduce 的 sent_bytes = size_bytes * 2 * (world_size-1)）就是你 benchmark 自己实现时要用的模式。

代码示例：数据并行（最简 DDP）

代码（Python）：

def data_parallelism_main(rank, world_size, data, num_layers, num_steps):
    setup(rank, world_size)
    batch_size = data.size(0)
    local_batch_size = int_divide(batch_size, world_size)     # 切分 batch
    data = data[rank*local_batch_size:(rank+1)*local_batch_size].to(f"cuda:{rank}")

    params = [get_init_params(num_dim, num_dim, rank) for _ in range(num_layers)]
    optimizer = torch.optim.AdamW(params, lr=1e-3)            # 每 rank 自己一份

    for step in range(num_steps):
        x = data
        for param in params:
            x = x @ param
            x = F.gelu(x)
        loss = x.square().mean()

        loss.backward()

        # 与单卡训练唯一的区别：
        for param in params:
            dist.all_reduce(tensor=param.grad, op=dist.ReduceOp.AVG, async_op=False)

        optimizer.step()
    cleanup()

代码做了什么： 每个 rank 用自己的 batch 切片算损失（各 rank 的 loss 不同），然后在 step 之前用 all-reduce(AVG) 把梯度跨 rank 平均——于是所有 rank 的参数演化完全一致。

实现深挖：

  • 为什么用 AVG 而不是 SUM：每个 rank 的梯度是其本地 batch 的均值；取平均才能保持对”全局 batch 梯度”的无偏估计。
  • 为什么说这是”朴素”DDP：每步一次阻塞 all-reduce，通信量 2×#params。讲义指出”下次：FSDP/ZeRO——用 all-gather 和 reduce-scatter 避免在显存里保存全部参数”。每 rank 显存：参数 + 梯度 + 优化器状态（bf16/AdamW 下 12+ 字节/参数）。
  • 为什么 MLP 有代表性：”MLP 是 Transformer 的计算瓶颈”——这些模式可以直接迁移。

与作业的联系：作业 2：实现带反向 hook 与异步通信的分布式数据并行（讲义的阻塞 all-reduce 是正确性基线，作业要求重叠）。这就是”DDP 必须做什么”的参考——其余都是工程细节。

代码示例：张量并行（前向）

代码（Python）：

def tensor_parallelism_main(rank, world_size, data, num_layers):
    setup(rank, world_size)
    data = data.to(f"cuda:{rank}")            # 所有 rank 都有完整 batch
    batch_size, num_dim = data.shape
    local_num_dim = int_divide(num_dim, world_size)   # 切宽度

    params = [get_init_params(num_dim, local_num_dim, rank) for _ in range(num_layers)]

    x = data
    for layer in range(num_layers):
        x = x @ params[layer]                 # 只用本 rank 的列切片
        x = F.gelu(x)

        activations = [torch.empty(batch_size, local_num_dim, device=f"cuda:{rank}")
                       for _ in range(world_size)]
        dist.all_gather(tensor_list=activations, tensor=x, async_op=False)  # 收集各切片
        x = torch.cat(activations, dim=1)     # 还原成完整宽度
    cleanup()

代码做了什么： 每个 rank 持有每层权重的一个列切片；完整 batch 依次通过各 rank 的切片，每层结束后把部分激活 all-gather 并拼接，得到下一层所需的完整宽度张量。

实现深挖：

  • 为什么每层都要 all-gather：下一层的矩阵乘需要完整激活；张量并行用”每层一次 all-reduce（8·b·s·h·(n−1)/n 每层）”换来权重的分片能力——这正是它必须依赖 NVLink 级带宽、通常只在节点内（≤8 卡）使用的原因。
  • 为什么反向是镜像关系：前向里 f 是恒等、g 是 all-reduce；反向里 f 变成 all-reduce、g 变成恒等——讲义把它作为”课后练习”，而第 8 讲幻灯片给出了模式：QKV/up-projection 做列并行，attention 输出/down-projection 做行并行。
  • 为什么显存线性扩展：每个 rank 只存每个权重矩阵的 1/world_size，因此参数显存随设备数线性下降。

与作业的联系：作业 2 的 FSDP 部分是这种显存分片思路的表亲；张量并行本身不在必修作业内（属第 8 讲内容），但你为 FSDP 写的 all-gather 机制用的是同一个原语。

代码示例：流水线并行（带 micro-batch）

代码（Python）：

def pipeline_parallelism_main(rank, world_size, data, num_layers, num_micro_batches):
    setup(rank, world_size)
    data = data.to(f"cuda:{rank}")
    batch_size, num_dim = data.shape

    local_num_layers = int_divide(num_layers, world_size)   # 切深度
    local_params = [get_init_params(num_dim, num_dim, rank) for _ in range(local_num_layers)]

    micro_batch_size = int_divide(batch_size, num_micro_batches)
    if rank == 0:
        micro_batches = data.chunk(chunks=num_micro_batches, dim=0)   # 数据源
    else:
        micro_batches = [torch.empty(micro_batch_size, num_dim, device=f"cuda:{rank}")
                         for _ in range(num_micro_batches)]

    for x in micro_batches:
        if rank - 1 >= 0:
            dist.recv(tensor=x, src=rank - 1)              # 从上一 stage 收激活
        for param in local_params:                          # 算我负责的层
            x = x @ param
            x = F.gelu(x)
        if rank + 1 < world_size:
            dist.send(tensor=x, dst=rank + 1)               # 传给下一个 stage
    cleanup()

代码做了什么： rank 0 持有数据；每个 rank 计算自己 stage 的层，并把激活发给下一个 rank。micro-batch 让 rank 1 可以在 rank 0 还在处理 micro-batch 1 时就开始算 micro-batch 0 的后半段——从而填满流水线气泡。

实现深挖：

  • 为什么要 micro-batch：没有它，同一时刻只有一个 rank 在工作（利用率 1/n）。有 m 个 micro-batch 时气泡比例约 (n−1)/m——”所以需要大 batch！”
  • 为什么用点对点 send/recv：流水线通信是”激活大小”且只发生在相邻 stage 之间——开销小到足以跑在慢速互联（节点间）上，这就是为什么流水线并行常跨机器、而张量并行留在节点内。
  • 为什么这里不做重叠：讲义明确不重叠通信与计算（”未处理：通过重叠通信/计算消除流水线气泡”）——工程实现里是带异步发送的 1F1B 调度。

与作业的联系：流水线并行不在必修作业中实现（作业 2 是 DDP → FSDP），但理解气泡数学能解释：为什么作业 2 的 FSDP 一定要做通信重叠，以及为什么讲义说”流水线能容忍慢互联，但需要花功夫减小气泡”。

关键要点

  1. 切分模型的方式很多：数据（batch）、张量/专家（宽度）、流水线（深度）、序列（长度）——各自的通信模式与硬件要求不同。
  2. 原语词表——broadcast/scatter/gather/reduce 及其 “all-“ 变体——是分布式训练的共同语言；all-reduce = reduce-scatter + all-gather，而这个分解正是 ZeRO/FSDP 得以成立的原因。
  3. 数据并行（DDP）：all-reduce 梯度、复制参数——简单，但通信 2×#params 且显存不随设备数下降。
  4. 张量并行：切分权重、每层 all-gather 激活——需要 NVLink；流水线并行：切分层、点对点传激活——能容忍慢网络但要付气泡代价。
  5. 需要最小化的是通信带宽（与第 5/6 讲最小化 HBM 带宽是同一个原则）；并用计算/通信重叠把它藏起来。

常见陷阱

  • 阻塞式集合操作串行化流水线：到处 async_op=False 会扼杀吞吐；要重叠并保存 handle。
  • 数据切分不均匀：int_divide 断言 a % b == 0——不均衡切分破坏 DDP 梯度平均的语义。
  • benchmark 忘记 dist.barrier()：各 rank 抢跑，计时里混入 straggler 偏斜。计时的集合操作两侧都要 barrier + synchronize。
  • all-gather 形状不匹配：tensor list 必须预先分配精确输出形状；all_gather_into_tensor 可免去手工管理。
  • 原地 all-reduce 的意外：输入张量会被覆盖——如需保留归约前的值要先拷贝。
  • 在慢链路上用重通信的集合操作：例如跨节点用以太网做逐层 all-reduce（张量并行）会完全支配运行时间；策略要与互联能力匹配。
  • 跨 rank 的随机种子控制：get_init_params 手动设了种子；不加控制会导致某些策略下初始化不一致（张量并行要求各权重切片语义一致，DDP 要求各 rank 初始参数相同）。可复现性很重要。

复习题

  1. 问： 为什么流水线并行能容忍慢互联，而张量并行不能？
    • 答： 流水线只在相邻 stage 之间传”激活”——每个 micro-batch O(b·s·h) 的点对点通信，与模型宽度无关。张量并行每一层都要 all-reduce 激活大小的张量——每层约 8·b·s·h·(n−1)/n，大致是流水线的 8 倍——所以它必须靠 NVLink 级带宽，否则通信将成为瓶颈。
  2. 问： 在 DDP 中，为什么梯度要取平均（AVG）而不是求和？
    • 答： 每个 rank 计算的梯度是其本地数据切片上的均值；跨 rank 求与会把有效学习率放大 world_size 倍。取平均才能复现全局 batch 的梯度。
  3. 问： 什么是流水线气泡？micro-batch 如何缩小它？
    • 答： 流水线一次通过的开头和结尾，各 stage 会在第一/最后一个 micro-batch 排空时闲置——闲置比例约 (n_stages−1)/m（m 为 micro-batch 数）。更多 micro-batch（更大的逻辑 batch）能填满流水线、摊薄气泡。