《CS336 Spring 2025 Lecture 8：Parallelism 2》学习笔记

(Updated: )
课程笔记
, ,

本文首发于知乎，现迁移至个人博客。

Lecture 8 主要是对 Lecture 7 涉及的并行计算基础理论进行具体的实现指引。

引言：从单 GPU 到多 GPU

上周我们讲了单张 GPU 内部的并行计算，本周聚焦多张 GPU 之间的并行计算。两种场景有一个共同主题：计算必须在 SM 内部的算术逻辑单元上执行，而所需的输入输出数据，存储位置可能相距较远——运气好的话数据就在 L1 缓存里，稍差一点在 HBM 里，而多 GPU、多节点训练场景下，你需要的数据甚至可能存放在另一张 GPU 上。

核心问题始终是：如何组织所有计算，以规避数据传输瓶颈？

上周通过算子融合与分块（Tiling），我们尽量减少对 HBM 的频繁读写——把数据加载到 L1 缓存后在本地完成计算，仅在必要时才写回 HBM。本周我们研究 GPU 与节点之间的通信：需要对模型参数和优化器状态进行复制与分片，而具体实现方式直接决定了训练开销。 整体硬件层次结构，从小而快到大而慢：

层级 存储 / 互联 备注
单节点、单 GPU L1 缓存 / 共享内存 速度最快，容量极小
单节点、单 GPU HBM H100 约 3.9 TB/s
单节点、多 GPU NVLink H100 约 900 GB/s
多节点、多 GPU NVSwitch 跨节点专用交换结构

减少数据传输的核心思想是通用的，只是在不同层级的具体机制不同——L1 缓存的工作方式和 NVSwitch 完全不一样。

本节课分为两部分： 第一部分：分布式通信基础组件

  • collective_operations()：集合通信操作的概念接口
  • torch_distributed()：NCCL 与 PyTorch 中的具体实现
  • benchmarking()：实测 NCCL 带宽

第二部分：分布式训练

  • data_parallelism()：沿批次维度切分
  • tensor_parallelism()：沿宽度（隐藏）维度切分
  • pipeline_parallelism()：沿深度（层数）维度切分

三种策略都基于深度 MLP 实现。MLP 是 Transformer 中的计算瓶颈而非注意力层，因此这个简单架构足以代表深度学习的典型负载。

Part 1：集合通信基础

1.1 集合通信操作

集合通信是分布式编程的基础原语，「集合」意味着涉及多个节点或设备，相比手动管理点对点通信，它提供了更高层次的抽象，是经过时间检验的可靠原语。这类技术至少从 20 世纪 80 年代的并行编程文献中就已出现。 两个基本术语：

  • world size：设备总数量，如 4
  • rank：设备编号，如 0、1、2、3（注意与线性代数中的「秩」无关）

六个核心操作：

操作 含义
broadcast 将某一 rank 上的数据发送到所有 rank
scatter 将不同数据分别发送到不同 rank
gather 将各 rank 上的不同数据汇集到某一 rank（scatter 的逆操作）
reduce 类似 gather，但对数据执行求和等运算而非简单拼接
all_gather 类似 gather，但结果发送到所有 rank
reduce_scatter 先对数据归约，再将结果的不同分片发送到不同 rank

如果是第一次看可能有点绕，帮助记忆的几个锚点：

  • reduce 表示执行求和、取最小/最大等可结合、可交换的运算
  • broadcast/scatter 与 gather 互为逆操作；
  • 前缀 all 表示目标是所有设备。

1.2 硬件实现

传统方案（家用/通用服务器）： 同一节点内 GPU 通过 PCI-E 总线通信，跨节点通过以太网。这种方案开销极大——GPU 间传输数据需经过主机内核、拷贝到缓冲区，再通过以太网发送。PCI-E 为通用场景设计（声卡、固态硬盘等设备都经由它连接），并非专为 GPU 通信优化。

现代数据中心方案： 既然明确要把多张 GPU 串联协同工作，英伟达提供了两种绕过传统链路的直连方案：

  • NVLink：同一节点内直接连接 GPU，绕过 CPU
  • NVSwitch：跨节点直接连接 GPU，绕过以太网

以 H100 为例：每张 GPU 配备 18 条第四代 NVLink，总带宽约 900 GB/s；而 HBM 内存带宽约 3.9 TB/s，也就是说 SM 访问 HBM 的速度仍比 NVLink 快约 4 倍。新一代 Blackwell 架构的互联带宽大约是当前的 2～3 倍，这些参数一直在演进。 用 nvidia-smi topo -m 可以查看集群实际的 GPU 连接拓扑。在我们的集群（8 张 GPU）上，每两张 GPU 之间都有 NV18 链路互联，同时还有 NIC（网卡）模块负责 GPU 通过 PCI-E 与 CPU 的通信——GPU 仍然需要和 CPU 通信，只是数据传输本身绕开了这条路径。

1.3 软件栈：NCCL 与 torch.distributed

英伟达的集合通信库 NCCL 把 all_reduce 这类高层操作转换为 GPU 之间的底层数据包。NCCL 初始化时会检测硬件拓扑（节点数量、交换结构、NVLink/PCIe 配置），自动优化 GPU 间的传输路径，再启动 CUDA 内核完成数据收发。对开发者来说，你只需要在代码层面声明「我需要这个张量同步到所有设备」，底层通信自动完成。

但 NCCL 对大多数开发者仍偏底层。PyTorch 的 torch.distributed 库在此之上提供了简洁的 Python 接口，可以直接对张量调用 all_gather、all_reduce 等函数。

torch.distributed 的另一个优点是支持多种后端适配不同硬件：GPU 使用 NCCL，CPU 使用 gloo。在没有 GPU 的笔记本上调试时，用 gloo 后端同样可以运行代码，保证程序逻辑正确。这也是高级原语的优势之一：可移植性远强于专用方案，性能最终仍取决于硬件，但至少逻辑上代码可以正常运行。

1.4 代码演示：集合通信操作

工具函数 spawn 基于 Python 多进程，启动 world_size 个进程执行同一段函数，每个进程对应一个 rank，运行完全相同的代码。因为讲课无法并行演示，我们只分步讲解单个 rank 的执行流程。

初始化与进程协调：

1
2
3
4
5
6
7
def setup(rank: int, world_size: int):
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "15623"
    if torch.cuda.is_available():
        dist.init_process_group("nccl", rank=rank, world_size=world_size)
    else:
        dist.init_process_group("gloo", rank=rank, world_size=world_size)

init_process_group 让所有进程连接到同一主机、互相感知彼此的存在。注意这一步只用于进程协调，实际数据通过 NCCL 传输，二者是分开的。 dist.barrier() 会阻塞当前进程，直到进程组内所有进程都执行到这一行，是多进程场景下设置同步点的标准手段。所有进程是异步运行的，barrier 保证步调对齐。

All-reduce

1
2
3
4
tensor = torch.tensor([0., 1, 2, 3], device=get_device(rank)) + rank
print(f"Rank {rank} [before all-reduce]: {tensor}")
dist.all_reduce(tensor=tensor, op=dist.ReduceOp.SUM, async_op=False)
print(f"Rank {rank} [after all-reduce]: {tensor}")

执行前：rank 0 持有 [0,1,2,3]，rank 1 持有 [1,2,3,4]，以此类推。 执行后：每一位都被所有 rank 的对应值求和覆盖——第一位 0+1+2+3=6，第二位 10，第三位 14，第四位 18。all_reduce 原地修改张量，每个 rank 得到完全相同的结果。async_op=False 表示同步执行，也可以用异步模式来实现通信与计算重叠。

Reduce-scatter

1
2
3
4
5
input = torch.arange(world_size, dtype=torch.float32, device=get_device(rank)) + rank
output = torch.empty(1, device=get_device(rank))
print(f"Rank {rank} [before reduce-scatter]: input={input}, output={output}")
dist.reduce_scatter_tensor(output=output, input=input, op=dist.ReduceOp.SUM, async_op=False)
print(f"Rank {rank} [after reduce-scatter]: input={input}, output={output}")

输入维度为 world_size（4），输出提前分配为标量。执行后，第一位的求和结果发给 rank 0，第二位发给 rank 1，以此类推。计算逻辑和 all_reduce 完全相同，只是输出被分散到了不同 rank，每个 rank 只拿到完整结果的一部分。 输入张量某一维度的长度必须等于 world_size，库会自动推断该维度对应哪个 rank 接收哪一分片，无需手动维护索引与 GPU 的映射关系。

All-gather（直接衔接 reduce_scatter 的输出）

1
2
3
4
5
input = output  # reduce-scatter 的输出
output = torch.empty(world_size, device=get_device(rank))
print(f"Rank {rank} [before all-gather]: input={input}, output={output}")
dist.all_gather_into_tensor(output_tensor=output, input_tensor=input, async_op=False)
print(f"Rank {rank} [after all-gather]: input={input}, output={output}")

执行后，所有 rank 都得到完整张量，与直接执行 all_reduce 的结果完全一致，直接验证了 all_reduce = reduce_scatter + all_gather。

1.5 性能基准测试

用 1 亿个 float32 元素的张量（约 400 MB），world_size=4，实测各操作的有效带宽。做基准测试必须先预热（执行一次操作完成内核加载），再用 torch.cuda.synchronize() 等待 CUDA 内核完成、dist.barrier() 等待所有进程就绪，确保计时准确。

All-reduce 带宽

1
2
3
4
size_bytes = tensor.element_size() * tensor.numel()
sent_bytes = size_bytes * 2 * (world_size - 1)  # 2x：先汇集做归约，再发回所有设备
total_duration = world_size * duration           # 4 个进程并行，等效单进程的 4 倍
bandwidth = sent_bytes / total_duration

2 倍系数的来源：all_reduce 需要先把所有数据汇集做归约，再把结果发回所有设备，每个 rank 既发送输入也接收输出，因此收发各算一次。 实测约 277 GB/s，H100 理论峰值约 900 GB/s。实际性能受张量大小、设备数量等因素影响，理论值和实测值之间的差距说明实测基准测试不可替代。

Reduce-scatter 带宽：

1
2
3
4
data_bytes = output.element_size() * output.numel()
sent_bytes = data_bytes * (world_size - 1)  # 没有 2x 系数
total_duration = world_size * duration
bandwidth = sent_bytes / total_duration

这里没有 2 倍系数：reduce_scatter 只需要把输入汇集一次做归约，scatter 只是把结果的不同分片分发到对应位置，本质是单向的归约操作。这也与等价关系相互印证——reduce_scatter 和 all_gather 各自都没有 2 倍系数，组合后正好对应 all_reduce 的 2 倍。 实测约 70 GB/s，与 all_reduce 差距较大。原因难以从理论上完全推导：all_reduce 的数据流量更大、NCCL 对其优化更充分，英伟达硬件还支持网络内计算（in-network computing）加速。NCCL 底层细节极多，这也正是我们需要实测而非依赖理论估算的原因。

Part 2：分布式训练

三种并行策略本质上是从不同维度切分模型或数据，对应不同的通信模式。

2.1 数据并行（Data Parallelism）

切分策略：沿批次维度切分数据，每个 rank 处理一个数据切片，同时持有完整的模型参数。

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

spawn(data_parallelism_main, world_size=4, data=data, num_layers=4, num_steps=1)

每个 rank 根据自身编号取对应的数据切片：

1
2
3
4
5
6
7
8
9
10
11
12
def data_parallelism_main(rank, world_size, data, num_layers, num_steps):
    setup(rank, world_size)
    batch_size    = data.size(0)                          # 128
    num_dim       = data.size(1)                          # 1024
    local_batch_size = int_divide(batch_size, world_size) # 32
    start_index   = rank * local_batch_size
    end_index     = start_index + local_batch_size
    data = data[start_index:end_index].to(get_device(rank))

    # 每个 rank 持有完整参数和独立的优化器状态
    params    = [get_init_params(num_dim, num_dim, rank) for _ in range(num_layers)]
    optimizer = torch.optim.AdamW(params, lr=1e-3)

训练循环与普通 SGD 完全相同，唯一的区别是在反向传播后、参数更新前插入一行梯度同步：

1
2
3
4
5
6
7
8
9
10
11
12
13
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()

       # 与标准训练的唯一区别：跨 rank 同步梯度
       for param in params:
           dist.all_reduce(tensor=param.grad, op=dist.ReduceOp.AVG, async_op=False)

       optimizer.step()

all_reduce 本身就是同步点，会阻塞所有进程直到全部完成，自然保证了训练步调一致。需要注意：如果某个 rank 少写了一次 all_reduce，程序就会卡死——其他进程会一直等待，永远不会超时。

执行结果：不同 rank 的 loss 不同（各自在不同数据切片上计算），但梯度经 all_reduce 平均后完全一致，参数更新后所有 rank 保持相同的参数。从 SGD 的视角看，流程没有任何变化，只是梯度被全局同步了。

这里体现了一个贯穿始终的权衡：优化器状态直接在本地更新，远比跨 GPU 传输更快，因此每个 rank 各自维护独立的优化器状态，只同步梯度。

关于批归一化的问题： 在语言模型中，我们很少使用批归一化，通常用层归一化（Layer Norm）。Layer Norm 只依赖当前样本的特征维度，不依赖批次维度，因此不同 rank 处理不同数据切片不会引入不一致。只要参数初始化和随机种子完全一致，结果就可以保证一致；GPU 上可能存在少量非确定性浮点误差，但影响很小。

2.2 张量并行（Tensor Parallelism）

切分策略：不切分数据，沿隐藏维度切分模型——每个 rank 只持有每一层参数的一部分，训练过程中需要频繁传输激活值。 张量并行的根本动机是：大模型无法放进单张 GPU，必须跨 GPU 分片存储参数。

1
2
3
4
5
6
7
8
9
10
11
spawn(tensor_parallelism_main, world_size=4, data=data, num_layers=4)

def tensor_parallelism_main(rank, world_size, data, num_layers):
    setup(rank, world_size)
    data = data.to(get_device(rank))
    batch_size    = data.size(0)                           # 128
    num_dim       = data.size(1)                           # 1024
    local_num_dim = int_divide(num_dim, world_size)        # 256

    # 参数矩阵：num_dim × local_num_dim（1024 × 256）
    params = [get_init_params(num_dim, local_num_dim, rank) for _ in range(num_layers)]

每层的前向传播：

1
2
3
4
5
6
7
8
9
10
11
12
x = data  # batch_size × num_dim
   for i in range(num_layers):
       x = x @ params[i]   # batch_size × local_num_dim（仅对参数切片做矩阵乘）
       x = F.gelu(x)

       # 用 all_gather 收集所有 rank 的局部激活值
       activations = [torch.empty(batch_size, local_num_dim, device=get_device(rank))
                      for _ in range(world_size)]
       dist.all_gather(tensor_list=activations, tensor=x, async_op=False)

       # 拼接得到完整激活张量：batch_size × num_dim
       x = torch.cat(activations, dim=1)

每一层计算后都需要执行 all_gather，把各 rank 的局部激活值拼接成完整张量再送入下一层。可以看到，张量并行在每一层都有通信开销，对设备互联带宽要求极高——这也是 Tatsu 上节课提到的：张量并行必须在 NVLink 这类高带宽互联上才能运行，否则大量的激活值传输会严重拖慢整体性能。 最终所有 rank 得到完整尺寸的激活张量，结果完全一致。反向传播逻辑较为繁琐，本节课略过。(不过结尾的时候 Percy 发现这节课还有时间时间，说 2026 可能会把反向传播加进来)

2.3 流水线并行（Pipeline Parallelism）

切分策略：按层数切分模型——所有 rank 都拿到完整的数据，每个 rank 只负责模型的某几层，层间通过点对点通信（send/recv）传递激活值。

1
2
3
4
5
6
7
8
9
spawn(pipeline_parallelism_main, world_size=2, data=data, num_layers=4, num_micro_batches=4)

def pipeline_parallelism_main(rank, world_size, data, num_layers, num_micro_batches):
    setup(rank, world_size)
    data = data.to(get_device(rank))
    batch_size      = data.size(0)                          # 128
    num_dim         = data.size(1)                          # 1024
    local_num_layers = int_divide(num_layers, world_size)   # 每 rank 负责 2 层
    local_params    = [get_init_params(num_dim, num_dim, rank) for _ in range(local_num_layers)]

4 层网络、2 个 rank，每个 rank 负责 2 层。 朴素流水线并行会产生流水线气泡（bubble），因为在某一 rank 计算时，其他 rank 只能等待。将总批次切分为微批次（micro-batches）可以减少气泡：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
micro_batch_size = int_divide(batch_size, num_micro_batches)  # 32
    if rank == 0:
        micro_batches = data.chunk(chunks=num_micro_batches, dim=0)
    else:
        micro_batches = [torch.empty(micro_batch_size, num_dim, device=get_device(rank))
                         for _ in range(num_micro_batches)]

    for x in micro_batches:
        if rank - 1 >= 0:
            dist.recv(tensor=x, src=rank - 1)     # 从上一个 rank 接收激活值
        for param in local_params:
            x = x @ param
            x = F.gelu(x)
        if rank + 1 < world_size:
            dist.send(tensor=x, dst=rank + 1)     # 发送给下一个 rank

逻辑直接：rank 0 处理输入数据的各个微批次，计算自己负责的层后发送给 rank 1；rank 1 等待接收、计算自己的层，以此类推。这里从点对点原语（send/recv）切换到集合通信原语（all_gather 等）也是可以的，但朴素实现用点对点更直观。 这个基础实现缺少几项重要的工程优化：

  • 通信与计算未重叠：目前 send/recv 是同步阻塞的。改为异步（isend 返回句柄，先发起发送，继续处理下一个微批次，最后统一等待）可以让通信和计算并行进行，消除等待开销。
  • 没有实现反向传播：加入反向传播后，还需要精细安排前向与反向步骤的穿插顺序，以进一步压缩流水线气泡。
  • 多流问题：如果同一 rank 向同一目标发送多次数据，数据按流保持顺序；不同 rank 之间的发送则可以任意时序，接收时需要指定源 rank 来区分。

关于「这是否类似事件驱动编程」的问题：事件驱动编程是注册回调、等待任意事件触发；流水线并行虽然也在等待前一个 rank 发来数据，但数据来源是固定的，整体仍是严格同步的编程范式。十几年前曾流行过异步训练——服务器下发数据，梯度就绪后自动累加，worker 异常也能容错，更接近事件驱动；但如今大规模深度学习训练，尽管规模极大，仍以同步范式为主。 最后一个 rank 得到完整前向传播的结果；加入反向传播时，从最后一个 rank 开始逐层向前回传梯度。

总结

三种并行策略对比：

并行方式 切分维度 每个 rank 持有 通信内容 通信操作
数据并行 批次（batch） 完整模型参数 梯度 all_reduce
张量并行 宽度（hidden dim） 部分参数 激活值 all_gather
流水线并行 深度（layers） 部分层参数 激活值 send / recv

除这三种之外，还有沿序列长度维度的序列并行。

一个贯穿两周的核心权衡：

  • 重新计算：节省内存，但消耗额外计算资源（激活重计算）
  • 存入内存：节省计算，但若数据在其他 GPU 上，必须承担通信开销
  • 多数情况下，系统瓶颈要么是通信，要么是显存

关于硬件演进： 硬件存储容量确实在增长，但模型规模永远会逼近硬件极限。从计算机系统诞生起，这种从快到慢的层次化存储结构就一直存在，未来也不会消失。

与工业界实现的差距

  • 本节课的示例缺少两类东西：一是更通用的模型支持（含注意力机制的 Transformer 而非深度 MLP）；二是通信与计算重叠的优化，这要求更复杂的状态管理。
  • 以 PyTorch FSDP 为例：要适配任意架构，就必须自动解析参数、追踪网络层位置，做大量簿记工作；而我们的 MLP 示例是手动指定了最简单的切分方式。Megatron-LM 和 PyTorch FSDP 的源码是理解工程实现复杂度的好参考。
  • Jax / TPU 生态是另一个值得了解的方向：Jax 面向计算图编程，你只需声明模型在各维度（模型维度、嵌入维度、注意力序列长度等）上如何切分，以及切分方式与 TPU 拓扑的映射关系，Jax 编译器会自动将其编译为底层通信原语。例如基于 Jax 开发的 Levanter 工具库，实现 FSDP 只需约 10 行代码：定义模型、指定切分维度，完成。这一抽象层级远高于手动编写集合通信代码。
  • 深度学习的计算图全程静态，从一开始就知道所有计算步骤。传统 GPU 诞生于需要大量分支的通用计算场景，并非专为这种静态数据流设计，因此专用芯片（如 Cerebras、Groq）有其独特优势——它们通过特殊制造工艺将内存直接集成在芯片上，几乎不需要片外数据搬运，代价是牺牲一定的通用灵活性。本课程坚持使用 PyTorch 从基础原语搭建，目的是让大家看清每一层的工作原理；在工业实践中，通常应该直接使用成熟的高层封装。