Lecture 8 · 并行 II：分布式训练实战

CS336: Language Modeling from Scratch · Stanford · Spring 2025
📅 Apr 24 · 💻 可执行讲义 · 讲者 Percy Liang

---

承上启下

上一讲（Lecture 7 · 并行 I：基础）：

• 为什么要并行：单 GPU 面临内存与算力双重瓶颈——训练 70B 模型需要约 $16\Psi \approx 1.12$ TB 的参数/梯度/优化器状态，远超单张 H100 的 80 GB HBM；而按 $C \approx 6ND$ 核算，单卡训练要花几百年。
• 集合通信原语（collective operations）：all-reduce / all-gather / reduce-scatter 等是多设备协同的抽象接口，恒等式 all-reduce = reduce-scatter + all-gather 是 ZeRO 系列的基石。
• 数据并行谱系：naive DP（复制完整模型）→ DDP（梯度同步）→ ZeRO 1/2/3（依次分片优化器状态/梯度/参数）→ FSDP（PyTorch 版 ZeRO-3）。
• 模型并行：张量并行（tensor parallelism，沿宽度切 MLP/注意力）、流水线并行（pipeline parallelism，沿深度切层）、序列并行（sequence parallelism，沿序列长度切逐点激活）。
• 3D 并行：数据 × 张量 × 流水线的组合，按「通信频率匹配链路带宽」的原则映射到硬件拓扑。

本讲（并行实战）：

• 上一讲给了理论地图——这一讲写出最小可运行代码，从原语开始亲手搭出每种并行。
• 从 torch.distributed 与 NCCL 的 hello world 开始，逐个实现数据并行 / 张量并行 / 流水线并行的核心循环。
• 用真实硬件（H100 / NVLink / 以太网）跑通信基准测试，把「带宽 900 GB/s」从参数表变成程序输出——延续 Lecture 6 · Kernel 与 Triton 的「先测量再优化」工作流。

---

1. 分布式通信的构建模块

1.1 集合通信原语：概念回顾

集合通信（collective operations） 是并行编程的经典抽象（1980 年代起源于 HPC 文献）——与其手写点对点通信（谁发给谁、谁等谁），不如一次性声明整个通信模式，让底层库（NCCL / MPI）根据硬件拓扑自动选择最优路由（ring、tree 等算法）。

核心术语：

• World size：参与设备数量 $W$（如 4 个 GPU）。
• Rank：设备编号 $0, 1, \dots, W-1$，每个 rank 对应一个进程。
• Reduce：满足结合律与交换律的汇聚运算（sum / min / max 等）——正因为满足这两条性质，归约的执行顺序才可以由通信库任意重排优化。

常用原语（本讲代码中出现的）：

操作	效果	典型用途
Broadcast	rank 0 的张量 → 所有 rank	分发初始参数
Scatter	rank 0 的张量切片 → 各 rank	分配数据批次
Gather	各 rank 张量 → rank 0 拼接	收集预测结果
Reduce	各 rank 张量求和 → rank 0	汇总损失
All-gather	各 rank 张量 → 所有 rank 拼接	FSDP 前向恢复完整参数
Reduce-scatter	各 rank 张量求和后切片分发	FSDP 反向梯度分片聚合
All-reduce	= reduce-scatter + all-gather	数据并行梯度同步

1.2 硬件拓扑：从 PCIe 到 NVSwitch

经典拓扑（家用/小规模）：

• 单节点内：GPU 通过 PCIe 总线连接（PCIe 7.0 × 16 lanes ≈ 242 GB/s），数据经常要绕道 CPU 内存；
• 跨节点：走以太网（约 200 MB/s 量级，比节点内慢约三个数量级）。

现代数据中心拓扑（H100 时代）：

• 单节点内：NVLink 直连 GPU，绕过 CPU（每张 H100 配 18 条 NVLink 4.0，总带宽 900 GB/s）；
• 跨节点 / 机柜级：通过 NVSwitch 与高速网络（InfiniBand / RoCE，或 GH200/GB200 的 NVLink Switch System）把 GPU 互联，尽量绕开传统以太网瓶颈（延迟更低、带宽更高）。

硬件拓扑图（现代数据中心）：

graph TB
    subgraph Node1["节点 1（8× H100）"]
        GPU0["GPU 0"] <-->|NVLink<br/>900GB/s| GPU1["GPU 1"]
        GPU1 <--> GPU2["GPU 2"]
        GPU2 <--> GPU3["GPU 3"]
        GPU0 <--> GPU3
    end

    subgraph Node2["节点 2（8× H100）"]
        GPU4["GPU 4"] <-->|NVLink<br/>900GB/s| GPU5["GPU 5"]
        GPU5 <--> GPU6["GPU 6"]
        GPU6 <--> GPU7["GPU 7"]
        GPU4 <--> GPU7
    end

    Node1 <-->|NVSwitch / IB<br/>跨节点互联| Node2

    style GPU0 fill:#c8e6c9,stroke:#2e7d32
    style GPU4 fill:#c8e6c9,stroke:#2e7d32

1.3 PyTorch 中的 torch.distributed：从 API 到 NCCL

分布式训练采用 SPMD（single program, multiple data）模式：同一份脚本被启动为 $W$ 个进程，每个进程绑一张 GPU，靠 rank 区分彼此的行为。

初始化分布式环境（每个 GPU 进程都要调用）：

import torch.distributed as dist
import os

def setup(rank: int, world_size: int):
    """初始化分布式进程组

    参数：
        rank: 当前进程的设备编号（0 到 world_size-1）
        world_size: 总设备数量
    """
    # 指定主节点地址（用于协调，实际数据走 NCCL）
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "15623"

    # 根据硬件选择后端：CUDA 用 NCCL，CPU 用 Gloo
    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)

def cleanup():
    """销毁进程组"""
    dist.destroy_process_group()

基本集合操作示例：

def all_reduce_example(rank: int, world_size: int):
    """all-reduce 示例：每个 rank 贡献一个数，求和后所有 rank 都得到结果"""
    setup(rank, world_size)

    # 每个 rank 创建自己的张量（值为 rank + 1）
    tensor = torch.tensor([rank + 1], dtype=torch.float32).cuda(rank)
    print(f"Rank {rank} [before]: {tensor.item()}")

    # all-reduce：对所有 rank 的张量求和，结果写回每个 rank 的 tensor
    dist.all_reduce(tensor=tensor, op=dist.ReduceOp.SUM, async_op=False)
    print(f"Rank {rank} [after]: {tensor.item()}")  # 输出：1+2+3+4 = 10

    cleanup()

# 输出示例（world_size=4）：
# Rank 0 [before]: 1.0  →  Rank 0 [after]: 10.0
# Rank 1 [before]: 2.0  →  Rank 1 [after]: 10.0
# Rank 2 [before]: 3.0  →  Rank 2 [after]: 10.0
# Rank 3 [before]: 4.0  →  Rank 3 [after]: 10.0

注意 all-reduce 是原地（in-place）操作：结果直接写回传入的 tensor。async_op=False 表示同步阻塞执行；设为 True 会返回一个 handle，允许通信与后续计算重叠（讲义的最小实现全部用同步版本，重叠优化留给工业级框架）。

all-reduce = reduce-scatter + all-gather（讲义用代码验证上一讲的恒等式）：

def verify_all_reduce_composition(rank: int, world_size: int):
    """验证 all-reduce 可分解为 reduce-scatter + all-gather"""
    setup(rank, world_size)

    # 创建输入矩阵（world_size × num_elements）
    num_elements = 4
    input_matrix = torch.arange(world_size * num_elements, dtype=torch.float32).cuda(rank)
    input_matrix = input_matrix.view(world_size, num_elements)

    # Step 1: reduce-scatter（对每行求和，第 i 行结果分发给 rank i）
    output_slice = torch.empty(num_elements).cuda(rank)
    dist.reduce_scatter_tensor(
        output=output_slice,
        input=input_matrix,
        op=dist.ReduceOp.SUM
    )
    print(f"Rank {rank} [after reduce-scatter]: {output_slice}")

    # Step 2: all-gather（收集所有 rank 的结果）
    final_output = torch.empty(world_size, num_elements).cuda(rank)
    dist.all_gather_into_tensor(
        output_tensor=final_output,
        input_tensor=output_slice
    )
    print(f"Rank {rank} [after all-gather]: {final_output}")

    cleanup()

1.4 通信带宽基准测试

测量 all-reduce 的实际带宽（代码来自 lecture_08.py）：

import time

def benchmark_all_reduce(rank: int, world_size: int, num_elements: int = 100 * 1024**2):
    """测量 all-reduce 的有效带宽

    参数：
        num_elements: 张量元素数量（默认 100M，约 400MB）
    """
    setup(rank, world_size)

    # 创建随机张量
    tensor = torch.randn(num_elements, device=f"cuda:{rank}")

    # Warmup：避免首次调用的初始化开销
    dist.all_reduce(tensor=tensor, op=dist.ReduceOp.SUM)
    torch.cuda.synchronize()  # 等待 GPU 计算完成
    dist.barrier()            # 等待所有进程就绪

    # 计时
    start = time.time()
    dist.all_reduce(tensor=tensor, op=dist.ReduceOp.SUM)
    torch.cuda.synchronize()
    dist.barrier()
    duration = time.time() - start

    # 计算带宽
    size_bytes = tensor.element_size() * tensor.numel()  # 数据量（字节）
    # all-reduce 的数据传输量 = size × 2 × (world_size - 1)
    # 因子 2：既要发送又要接收；因子 (world_size - 1)：与其他 rank 通信
    sent_bytes = size_bytes * 2 * (world_size - 1)
    total_duration = world_size * duration  # 总时间 = 单次耗时 × 设备数
    bandwidth_gbps = (sent_bytes / total_duration) / (1024**3)

    print(f"Rank {rank}: all-reduce 耗时 {duration*1000:.2f} ms, "
          f"测得带宽 {bandwidth_gbps:.1f} GB/s")

    cleanup()

# 典型输出（4× H100，单节点 NVLink）：
# Rank 0: all-reduce 耗时 8.3 ms, 测得带宽 289 GB/s
# Rank 1: all-reduce 耗时 8.3 ms, 测得带宽 289 GB/s
# ...
# （理论峰值 900 GB/s，实测约 30%，受算法与拓扑限制）

代码中的带宽公式即 NCCL 官方基准（nccl-tests）所谓的「算法带宽/总线带宽」约定。对大小 $s$ 字节的张量、$W$ 个设备，all-reduce 的有效带宽为：

$$B_{\text{all-reduce}} = \frac{2(W-1)\,s}{W \cdot t} = \frac{2\,\frac{W-1}{W}\,s}{t}$$

其中 $t$ 是单次 all-reduce 耗时。分子正是上一讲推导的 ring all-reduce 每设备收发字节数（reduce-scatter 与 all-gather 各贡献 $\frac{W-1}{W}s$，故有因子 2）；类似地 reduce-scatter 单独测时分子只有 $\frac{W-1}{W}s$（无因子 2）。这样定义的带宽可以直接与 NVLink 标称值对比。

---

2. 数据并行：切分 batch 维度

2.1 核心思想

数据并行（Data Parallelism, DP/DDP） 是最直观的并行策略：

• 参数复制：每个 GPU 持有模型的完整副本；
• 数据切分：把 batch 切成 $W$ 份，每个 GPU 处理一部分；
• 梯度同步：反向传播后，用 all-reduce 对所有 GPU 的梯度求平均，保证各副本走完全相同的优化轨迹。

graph LR
    subgraph "Rank 0"
        D0["数据<br/>batch[0:32]"] --> M0["模型副本<br/>θ"]
        M0 --> L0["loss₀"]
        L0 --> G0["∇₀"]
    end

    subgraph "Rank 1"
        D1["数据<br/>batch[32:64]"] --> M1["模型副本<br/>θ"]
        M1 --> L1["loss₁"]
        L1 --> G1["∇₁"]
    end

    subgraph "Rank 2"
        D2["数据<br/>batch[64:96]"] --> M2["模型副本<br/>θ"]
        M2 --> L2["loss₂"]
        L2 --> G2["∇₂"]
    end

    subgraph "Rank 3"
        D3["数据<br/>batch[96:128]"] --> M3["模型副本<br/>θ"]
        M3 --> L3["loss₃"]
        L3 --> G3["∇₃"]
    end

    G0 --> AR["all-reduce<br/>求平均梯度"]
    G1 --> AR
    G2 --> AR
    G3 --> AR

    AR --> U["更新参数<br/>θ ← θ - lr·∇"]

    style AR fill:#ffe0b2,stroke:#e65100
    style U fill:#c8e6c9,stroke:#2e7d32

同步后的梯度为 $\bar g = \frac{1}{W}\sum_{w=1}^{W} g_w$，数学上等价于在完整全局 batch 上做一次大 batch 梯度下降——并行只是改变了计算的物理分布，不改变优化语义。

2.2 最小实现：深度 MLP 的数据并行

import torch
import torch.nn.functional as F
import torch.distributed as dist

def data_parallelism_train(rank: int, world_size: int, data: torch.Tensor, 
                           num_layers: int = 4, num_steps: int = 10):
    """数据并行训练深度 MLP

    参数：
        rank: 设备编号
        world_size: 总设备数
        data: 完整训练数据（shape: [batch_size, num_dim]）
        num_layers: MLP 层数
        num_steps: 训练步数
    """
    setup(rank, world_size)

    # 1. 切分数据：每个 rank 只取自己的那部分
    batch_size, num_dim = data.shape
    local_batch_size = batch_size // world_size
    start_idx = rank * local_batch_size
    end_idx = start_idx + local_batch_size
    local_data = data[start_idx:end_idx].cuda(rank)

    # 2. 初始化模型参数（每个 rank 持有完整副本）
    params = [
        torch.randn(num_dim, num_dim, requires_grad=True, device=f"cuda:{rank}")
        for _ in range(num_layers)
    ]
    optimizer = torch.optim.AdamW(params, lr=1e-3)

    # 3. 训练循环
    for step in range(num_steps):
        # 前向传播（只在本地数据上）
        x = local_data
        for W in params:
            x = x @ W
            x = F.gelu(x)
        loss = x.square().mean()  # 简单损失：输出的 L2 范数

        # 反向传播（计算本地梯度）
        loss.backward()

        # ★ 关键步骤：同步梯度（用平均而非求和）
        for W in params:
            dist.all_reduce(W.grad, op=dist.ReduceOp.AVG)

        # 更新参数（所有 rank 的参数保持一致）
        optimizer.step()
        optimizer.zero_grad()

        if rank == 0:  # 只让 rank 0 打印
            print(f"Step {step}: loss = {loss.item():.4f}")

    cleanup()

# 启动多进程（使用 torch.multiprocessing）
if __name__ == "__main__":
    world_size = 4
    data = torch.randn(128, 1024)  # 128 个样本，每个 1024 维
    torch.multiprocessing.spawn(
        data_parallelism_train,
        args=(world_size, data),
        nprocs=world_size,
        join=True
    )

2.3 关键观察

运行输出解析（来自讲义）：

[data_parallelism] Rank 0: step = 0, loss = 341.2, params = [mean=0.01, std=0.58]
[data_parallelism] Rank 1: step = 0, loss = 338.9, params = [mean=0.01, std=0.58]
[data_parallelism] Rank 2: step = 0, loss = 342.1, params = [mean=0.01, std=0.58]
[data_parallelism] Rank 3: step = 0, loss = 339.7, params = [mean=0.01, std=0.58]

• loss 不同：每个 rank 处理的数据不同，本地损失自然不同——这是数据并行工作正常的标志，不是 bug；
• 参数相同：梯度同步后所有 rank 的参数统计量一致（mean/std），说明各副本确实走在同一条优化轨迹上；
• 扩展性：通信量 $=O(\text{模型大小})$、与 batch size 无关，所以 batch 越大、每步计算越多，通信占比越低。

通信开销分析：

• 每步一次梯度 all-reduce，同步的数据量等于梯度大小 $G$（70B 模型 BF16 梯度即 $G = 140$ GB）；
• 由第 1.4 节公式，每卡实际收发约 $2\frac{W-1}{W}G \approx 2G = 280$ GB，按 300 GB/s 有效带宽约需 0.9 秒；
• 若前向 + 反向计算需数秒，通信占比约两到三成——不可忽略，因此工业实现（DDP）必须把 all-reduce 与反向传播重叠：某层梯度一算完就异步发起通信，藏在更早层的反向计算后面。

---

3. 张量并行：切分宽度维度

3.1 核心思想

张量并行（Tensor Parallelism, TP） 将模型的每一层按宽度切分——每个 GPU 只持有参数矩阵的一部分列。

以 MLP 为例（$y = \mathrm{GeLU}(xW)$，$W$ 形状 $[d, d]$）：

graph LR
    subgraph "单 GPU（完整）"
        X1["x<br/>[B, d]"] --> W1["W<br/>[d, d]"]
        W1 --> Y1["y<br/>[B, d]"]
    end

    subgraph "张量并行（4 GPU）"
        X2["x<br/>[B, d]<br/>（复制到所有 GPU）"] --> W0["W₀<br/>[d, d/4]<br/>GPU 0"]
        X2 --> W1p["W₁<br/>[d, d/4]<br/>GPU 1"]
        X2 --> W2["W₂<br/>[d, d/4]<br/>GPU 2"]
        X2 --> W3["W₃<br/>[d, d/4]<br/>GPU 3"]

        W0 --> Y0["y₀<br/>[B, d/4]"]
        W1p --> Y1p["y₁<br/>[B, d/4]"]
        W2 --> Y2["y₂<br/>[B, d/4]"]
        W3 --> Y3["y₃<br/>[B, d/4]"]

        Y0 --> AG["all-gather"]
        Y1p --> AG
        Y2 --> AG
        Y3 --> AG

        AG --> YF["y<br/>[B, d]<br/>（拼接）"]
    end

    style AG fill:#ffe0b2,stroke:#e65100
    style YF fill:#c8e6c9,stroke:#2e7d32

关键步骤：

1. 参数切分：$W$ 按列切成 $W_{\text{world}}$ 份，rank $i$ 持有列块 $W_i \in \mathbb{R}^{d \times d/4}$（对应 W[:, i*d/4:(i+1)*d/4]）；
2. 激活复制：输入 $x$ 在所有 GPU 上保持完整副本；
3. 局部计算：每个 GPU 计算 $y_i = \mathrm{GeLU}(xW_i)$——GeLU 逐元素作用，对列分片可独立施加；
4. all-gather：收集各 GPU 的输出列块，拼接成完整的 $y = [y_0 \mid y_1 \mid y_2 \mid y_3]$，作为下一层的输入。

3.2 最小实现：深度 MLP 的张量并行

def tensor_parallelism_train(rank: int, world_size: int, data: torch.Tensor, num_layers: int = 4):
    """张量并行训练深度 MLP

    参数：
        rank: 设备编号
        world_size: 总设备数
        data: 完整训练数据（所有 rank 使用相同数据）
        num_layers: MLP 层数
    """
    setup(rank, world_size)

    # 1. 每个 rank 持有完整数据（张量并行不切分 batch）
    batch_size, num_dim = data.shape
    local_data = data.cuda(rank)

    # 2. 参数切分：每个 rank 只持有参数矩阵的 1/world_size 列
    local_dim = num_dim // world_size
    params = [
        torch.randn(num_dim, local_dim, requires_grad=True, device=f"cuda:{rank}")
        for _ in range(num_layers)
    ]

    # 3. 前向传播
    x = local_data
    for W_local in params:
        # 3.1 局部矩阵乘法：x @ W_local → [batch_size, local_dim]
        x_local = x @ W_local
        x_local = F.gelu(x_local)

        # 3.2 all-gather：收集所有 GPU 的输出
        x_gathered = [torch.empty_like(x_local) for _ in range(world_size)]
        dist.all_gather(tensor_list=x_gathered, tensor=x_local)

        # 3.3 拼接：恢复完整维度 [batch_size, num_dim]
        x = torch.cat(x_gathered, dim=1)

    loss = x.square().mean()
    print(f"Rank {rank}: loss = {loss.item():.4f}, "
          f"参数形状 = {params[0].shape}")  # 输出 [1024, 256]（若 world_size=4, num_dim=1024）

    # 反向传播（省略：all-gather 不自动传梯度，需自定义 autograd 函数，讲义留作练习）
    cleanup()

3.3 关键观察

运行输出（world_size=4, num_dim=1024）：

Rank 0: loss = 342.1, 参数形状 = torch.Size([1024, 256])
Rank 1: loss = 342.1, 参数形状 = torch.Size([1024, 256])
Rank 2: loss = 342.1, 参数形状 = torch.Size([1024, 256])
Rank 3: loss = 342.1, 参数形状 = torch.Size([1024, 256])

• loss 相同：所有 rank 处理相同数据，且每层输出经 all-gather 拼回完整张量，前向结果与单卡完全一致；
• 参数更小：每个 GPU 只存 $1/4$ 参数（1024→256 列）——这是张量并行的内存收益；
• 通信密集：每层都要一次 all-gather，传输的是激活，大小 = batch_size × num_dim。

通信开销分析：

• 每层 all-gather 激活的数据量为 $b_{\text{tok}} \cdot h \cdot 2$ 字节（$b_{\text{tok}}$ 为本次前向的 token 数，$h$ 为 hidden dim，BF16 每元素 2 字节）；
• 例：每步 100 万 token、$h = 4096$、BF16，则单层单向就是 $10^6 \times 4096 \times 2 \approx 8$ GB；32 层前向合计约 256 GB，反向还要再来一轮——按激活算的通信量远超数据并行按梯度算的通信量；
• 这就是张量并行「每层收税」的含义：通信频率是每层一次（而非每步一次），且不容易与计算重叠。

---

4. 流水线并行：切分深度维度

4.1 核心思想

流水线并行（Pipeline Parallelism, PP） 将模型按层切分——每个 GPU 负责连续几层（一个 stage），激活在 GPU 间依次流动。

graph LR
    subgraph "GPU 0（层 0-7）"
        L0["Layer 0-7"] --> A0["激活₀"]
    end

    subgraph "GPU 1（层 8-15）"
        A0 --> L1["Layer 8-15"] --> A1["激活₁"]
    end

    subgraph "GPU 2（层 16-23）"
        A1 --> L2["Layer 16-23"] --> A2["激活₂"]
    end

    subgraph "GPU 3（层 24-31）"
        A2 --> L3["Layer 24-31"] --> A3["输出"]
    end

    style A0 fill:#ffe0b2,stroke:#e65100
    style A1 fill:#ffe0b2,stroke:#e65100
    style A2 fill:#ffe0b2,stroke:#e65100

挑战：流水线气泡（pipeline bubble）——朴素实现下游阶段等上游、上游算完等下游，大量时间空转：

gantt
    title 朴素流水线（单 batch）：GPU 利用率低
    dateFormat X
    axisFormat %L

    section GPU 0
    层 0-7 前向 :a1, 0, 1
    等待         :a2, 1, 3
    层 0-7 反向 :a3, 3, 1

    section GPU 1
    等待         :b1, 0, 1
    层 8-15 前向:b2, 1, 1
    等待         :b3, 2, 1
    层 8-15 反向:b4, 3, 1

    section GPU 2
    等待         :c1, 0, 2
    层 16-23 前向:c2, 2, 1
    层 16-23 反向:c3, 3, 1

    section GPU 3
    等待         :d1, 0, 3
    层 24-31 前向:d2, 3, 1

解决方案：微批次（micro-batching）——把 batch 切成 $m$ 个小批次流水线式填充，让各阶段错峰同时工作：

gantt
    title 微批次流水线（4 个 micro-batch）：GPU 利用率高
    dateFormat X
    axisFormat %L

    section GPU 0
    MB0 前向:a0, 0, 1
    MB1 前向:a1, 1, 1
    MB2 前向:a2, 2, 1
    MB3 前向:a3, 3, 1
    MB3 反向:a4, 7, 1
    MB2 反向:a5, 8, 1
    MB1 反向:a6, 9, 1
    MB0 反向:a7, 10, 1

    section GPU 1
    等待    :b0, 0, 1
    MB0 前向:b1, 1, 1
    MB1 前向:b2, 2, 1
    MB2 前向:b3, 3, 1
    MB3 前向:b4, 4, 1
    MB3 反向:b5, 6, 1
    MB2 反向:b6, 7, 1

    section GPU 2
    等待    :c0, 0, 2
    MB0 前向:c1, 2, 1
    MB1 前向:c2, 3, 1
    MB2 前向:c3, 4, 1
    MB3 前向:c4, 5, 1
    MB3 反向:c5, 6, 1

    section GPU 3
    等待    :d0, 0, 3
    MB0 前向:d1, 3, 1
    MB1 前向:d2, 4, 1
    MB2 前向:d3, 5, 1
    MB3 前向:d4, 6, 1

4.2 最小实现：深度 MLP 的流水线并行

def pipeline_parallelism_train(rank: int, world_size: int, data: torch.Tensor, 
                               num_layers: int = 4, num_micro_batches: int = 4):
    """流水线并行训练深度 MLP

    参数：
        rank: 设备编号（对应流水线阶段）
        world_size: 总设备数
        data: 完整训练数据
        num_layers: 总层数
        num_micro_batches: 微批次数量（越大气泡越小）
    """
    setup(rank, world_size)

    batch_size, num_dim = data.shape
    micro_batch_size = batch_size // num_micro_batches

    # 1. 层切分：每个 rank 负责 num_layers / world_size 层
    local_num_layers = num_layers // world_size
    params = [
        torch.randn(num_dim, num_dim, requires_grad=True, device=f"cuda:{rank}")
        for _ in range(local_num_layers)
    ]

    # 2. 准备微批次
    if rank == 0:
        # Rank 0 持有原始数据，切分成微批次
        micro_batches = data.chunk(num_micro_batches, dim=0)
        micro_batches = [mb.cuda(rank) for mb in micro_batches]
    else:
        # 其他 rank 分配接收缓冲区
        micro_batches = [
            torch.empty(micro_batch_size, num_dim, device=f"cuda:{rank}")
            for _ in range(num_micro_batches)
        ]

    # 3. 前向传播（流水线式）
    for i, x in enumerate(micro_batches):
        # 3.1 接收上一阶段的激活
        if rank > 0:
            dist.recv(tensor=x, src=rank - 1)

        # 3.2 计算本阶段的层
        for W in params:
            x = x @ W
            x = F.gelu(x)

        # 3.3 发送到下一阶段
        if rank < world_size - 1:
            dist.send(tensor=x, dst=rank + 1)
            print(f"Rank {rank}: 发送 micro-batch {i} 到 rank {rank+1}")

    # 反向传播（省略，需逆序处理微批次）
    cleanup()

# 输出示例（world_size=2, num_micro_batches=4）：
# Rank 0: 发送 micro-batch 0 到 rank 1
# Rank 0: 发送 micro-batch 1 到 rank 1
# Rank 0: 发送 micro-batch 2 到 rank 1
# Rank 0: 发送 micro-batch 3 到 rank 1

注意这是三种并行里唯一用点对点通信（dist.send / dist.recv）而非 collective 的：通信天然只发生在相邻阶段之间。讲义明确注明这个最小版没有做通信与计算的重叠，气泡原样存在——消除气泡正是工业级调度器的核心工作。

4.3 关键观察与局限

优点：

• 参数内存线性下降：每个 GPU 只存 $1/p$ 的层；
• 通信量小：每个 micro-batch 只在相邻阶段间传一次层间激活（大小 = micro_batch_size × hidden_dim），是三种并行中通信最省的，因此最适合跨节点低带宽链路；
• 实现直观：点对点 send/recv，无需全局同步。

缺点：

• 气泡开销：即使用微批次，仍有约 $(p-1)/m$ 的时间浪费；
• 激活内存：朴素调度（GPipe 式「全前向再全反向」）要同时保存全部 $m$ 个微批次的激活等待反向；
• 依赖大 batch：必须有足够多 micro-batch 才能摊平气泡，小 batch 场景（如 RL 微调、推理）流水线很吃亏。

实战优化（GPipe / PipeDream / Megatron-LM）：

• 1F1B 调度（one-forward-one-backward）：稳态时每做一个前向就接一个反向，把在途激活从 $m$ 份压到 $p$ 份量级，气泡不变但内存峰值大降；
• 交错流水线（interleaved schedule）：每卡持有多段不连续的层，进一步缩小气泡（代价是更多通信）；
• 激活检查点（checkpointing）：只存边界激活，反向时重算内部激活，用约 1/3 的额外计算换内存。

---

5. 3D 并行：组合三大策略

5.1 为什么需要组合

单一策略的瓶颈：

• 数据并行：每步 all-reduce 梯度，通信量 $\propto$ 模型大小（70B 模型约 140 GB/步），且不切分任何内存；
• 张量并行：每层 all-gather/all-reduce 激活，通信量 $\propto$ 层数 × token 数 × hidden dim，出了 NVLink 就撑不住；
• 流水线并行：通信最省但有气泡，需要大 batch 摊平。

组合策略（3D Parallelism）：同时使用三种并行，让每种并行的通信落在它能承受的链路上：

graph TB
    subgraph DP0["数据并行组 0"]
        subgraph PS00["流水线阶段 0"]
            TP00["张量并行<br/>GPU 0-3"]
        end
        subgraph PS01["流水线阶段 1"]
            TP01["张量并行<br/>GPU 4-7"]
        end
    end

    subgraph DP1["数据并行组 1"]
        subgraph PS10["流水线阶段 0"]
            TP10["张量并行<br/>GPU 8-11"]
        end
        subgraph PS11["流水线阶段 1"]
            TP11["张量并行<br/>GPU 12-15"]
        end
    end

    TP00 -.数据并行<br/>梯度同步.-> TP10
    TP01 -.数据并行<br/>梯度同步.-> TP11

    style TP00 fill:#c8e6c9,stroke:#2e7d32
    style TP10 fill:#c8e6c9,stroke:#2e7d32

实测规模的参考配置（Megatron 论文 Narayanan et al. 2021 中 GPT-3 175B 规模的基准）：

• 张量并行：8 路（恰好占满一个 DGX 节点的 NVLink 域）；
• 流水线并行：16 路（跨节点，通信频率低）；
• 数据并行：12 路（最外层，梯度同步与计算重叠）；
• 总计：$8 \times 16 \times 12 = 1536$ 张 A100，实现约 52% MFU。

三个并行度的乘积等于总卡数：$W = W_{\text{TP}} \times W_{\text{PP}} \times W_{\text{DP}}$。

5.2 通信模式对比

并行策略	通信范围	通信量/步	通信频率	适用场景
数据并行	DP 组内全局 all-reduce	O(参数量)	每步一次	大 batch，可与计算重叠
张量并行	单节点内 all-gather/all-reduce	O(激活) × 层数	每层多次	节点内 NVLink 高带宽
流水线并行	相邻阶段点对点 send/recv	O(激活)	每微批次一次	跨节点，层数多

---

6. 实战工具链

6.1 PyTorch 原生支持

DDP（DistributedDataParallel）：数据并行的高级封装

from torch.nn.parallel import DistributedDataParallel as DDP

model = MyModel().cuda(rank)
ddp_model = DDP(model, device_ids=[rank])

# DDP 自动在 backward() 后插入 all-reduce
output = ddp_model(input)
loss = criterion(output, target)
loss.backward()  # ← 梯度同步在这里自动完成
optimizer.step()

相比第 2 节的手写版本，DDP 额外做了梯度分桶（把小梯度攒成大桶减少通信次数）和通信计算重叠（桶满即异步 all-reduce）。

FSDP（Fully Sharded Data Parallel）：ZeRO-3 的 PyTorch 实现

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

model = MyModel()
fsdp_model = FSDP(
    model,
    sharding_strategy=ShardingStrategy.FULL_SHARD,  # 参数 + 梯度 + 优化器状态全切分
    cpu_offload=CPUOffload(offload_params=True),    # 可选：卸载到 CPU
)

# 使用方式与 DDP 相同，但内存占用大幅降低
output = fsdp_model(input)
loss.backward()
optimizer.step()

FSDP 逐模块 all-gather 参数、用完释放，并预取下一模块的参数以掩盖通信——原理即 Lecture 7 · 并行 I：基础 第 4.3 节的 ZeRO-3 流程。讲义有意不直接教这些高层 API 的调参，而是先用裸原语搭一遍，让你知道封装之下发生了什么。

6.2 第三方库

Megatron-LM（NVIDIA）：

• 工业级 3D/4D 并行实现（张量 + 流水线 + 数据 + 序列/上下文），针对 Transformer 结构手工安排每一次通信的位置；
• 用于训练 GPT、BERT、T5 及众多开源大模型。

DeepSpeed（Microsoft）：

• ZeRO 系列优化的参考实现（ZeRO-1/2/3/Infinity），支持混合精度、梯度累积、CPU/NVMe 卸载；
• 一行代码集成：deepspeed.initialize(model, optimizer, ...)。

JAX + TPU：

• 自动分片（GSPMD）：只需标注张量的分片策略，编译器自动插入所有通信——与 PyTorch 手工编排形成鲜明对比；
• Levanter / EasyLM 等高层框架基于此构建（讲义点名 Levanter 作为「声明式并行」的代表）。

---

总结

mindmap
  root((并行 II<br/>分布式训练实战))
    构建模块
      集合通信原语
        all-reduce / all-gather
        reduce-scatter
      torch.distributed
        NCCL 后端
        setup / cleanup
      通信基准测试
        测量带宽
        理解瓶颈
    数据并行
      切分 batch
      all-reduce 梯度
      DDP / FSDP
    张量并行
      切分宽度
      all-gather 激活
      Megatron-LM
    流水线并行
      切分深度
      微批次流水线
      气泡开销
    3D 并行
      组合三大策略
      匹配硬件拓扑
      Megatron 175B 配置

关键要点：

1. 集合通信原语是分布式训练的基石——all-reduce = reduce-scatter + all-gather（讲义用代码验证）；理解每个原语的数据流向，才能读懂并设计并行策略。
2. 三大并行策略各有权衡——数据并行最简单但每步付「模型大小」的通信税且不省内存；张量并行切参数但每层同步激活；流水线并行通信最省但有 $(p-1)/m$ 的气泡。
3. 硬件拓扑决定并行选择——节点内 NVLink（900 GB/s）养得起张量并行的高频通信；跨节点走流水线并行的低频点对点；数据并行放最外层并与反向传播重叠。
4. 通信是新瓶颈——NVLink 比 HBM 慢约 4 倍；70B 模型一次梯度 all-reduce 每卡收发约 280 GB、按实测 300 GB/s 需近 1 秒；优化的两条路永远是「减少通信量」与「重叠通信与计算」。
5. 从最小代码到工业实现——讲义的裸原语版本刻意省略了：通信计算重叠（async_op=True + CUDA 流）、梯度分桶、1F1B/交错调度、混合精度与梯度累积、激活检查点、故障恢复。这些正是 Megatron/DeepSpeed 的增量价值，也是 CS336 作业 2 要亲手补上的部分。

下一讲预告（Lecture 9 · 缩放定律 I：基础）：

从「怎么并行训练」转向「训练多大的模型」——缩放定律（scaling laws）告诉我们：

• 幂律（power laws）：loss 如何随模型大小 $N$、数据量 $D$、计算量 $C$ 变化；
• Kaplan 2020：OpenAI 的原始缩放定律（”bigger is better”）；
• Chinchilla 2022：DeepMind 的修正——模型与数据需同步增长（约 20 tokens/参数）；
• 实战意义：给定计算预算，如何分配 $N$ 和 $D$ 以最大化性能。

---

复习自测

---

参考资料

论文：
- Megatron-LM（NVIDIA, 2019）：Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
- 提出张量并行（Tensor Parallelism）在 Transformer 中的实现
- GPipe（Google, 2019）：GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism
- 流水线并行 + 微批次 + 梯度累积
- ZeRO（Microsoft, 2020）：ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
- 数据并行的内存优化（分片优化器状态、梯度、参数）
- PyTorch FSDP（Meta, 2021）：Fully Sharded Data Parallel: faster AI training with fewer GPUs
- ZeRO-3 的 PyTorch 原生实现
- 3D 并行实测（NVIDIA/Stanford/MSR, 2021）：Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM
- TP × PP × DP 组合的系统性实验与 175B 规模配置

工具与教程：
- PyTorch Distributed 教程：https://pytorch.org/tutorials/beginner/dist_overview.html
- NCCL 官方文档：https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/
- NCCL 性能测试工具：https://github.com/NVIDIA/nccl-tests
- Megatron-LM 代码库：https://github.com/NVIDIA/Megatron-LM
- DeepSpeed 代码库：https://github.com/microsoft/DeepSpeed
- Stas Bekman 的 ML 工程笔记（通信基准测试代码）：https://github.com/stas00/ml-engineering

课程相关：
- Lecture 7 · 并行 I：基础：概念地图、集合通信原语、ZeRO / FSDP 原理
- Lecture 9 · 缩放定律 I：基础：Kaplan 2020、Chinchilla 2022、20 tokens/参数的经验法则
- CS336 作业 2：实现单机多卡 + 多机分布式训练（DDP / FSDP）

延伸阅读：
- Transformer 推理优化（Lecture 10 · 推理）：推理侧的并行与内存问题——KV cache 分片、投机解码、量化
- 数据并行 vs 模型并行的历史演进：从 AlexNet（2012）把卷积核切到两张 GTX 580 的朴素模型并行，到 GPT-3（2020）时代成体系的 3D 并行
- 通信与计算重叠：异步 all-reduce（async_op=True）+ CUDA 流调度，是从「能跑」到「跑满」的关键一步

---