# Lecture 8 · 并行 II：分布式训练实战 > **CS336: Language Modeling from Scratch** · Stanford · Spring 2025 > 📅 Apr 24 · 💻 [可执行讲义](https://github.com/stanford-cs336/spring2025-lectures/blob/main/lecture_08.py) · 讲者 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]] 的「先测量再优化」工作流。 > [!abstract] 一句话总览 > 上一讲画了作战地图，这一讲拿起工兵铲挖战壕——用不到 200 行 PyTorch 代码，把三大并行策略从图纸变成可执行的 Python 函数。 --- ## 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 | 数据并行梯度同步 | > [!tip] 记忆技巧 > - **Reduce** = 执行结合/交换运算（sum / min / max）； > - **Scatter / Gather** 互为逆操作（切片分发 vs 拼接收集）； > - **All-** 前缀 = 结果送达所有设备（而非只汇聚到 rank 0）。 > 用这三条规则可以从名字直接反推出每个原语的数据流向，不用死记七个操作。 ### 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 互联，尽量绕开传统以太网瓶颈（延迟更低、带宽更高）。 > [!warning] 关键对比 > H100 的 NVLink 带宽（900 GB/s）仍比自家 HBM 带宽（3.9 TB/s，见 [[Lecture 5 · GPU]]）慢约 **4 倍**，跨节点又比 NVLink 慢一个数量级——存储层级每往外一层就掉一档带宽。GPU 间通信必然是瓶颈，这就是为什么要精心设计通信模式，而不是「能通就行」。 **硬件拓扑图（现代数据中心）**： ```mermaid graph TB subgraph Node1["节点 1（8× H100）"] GPU0["GPU 0"] <-->|NVLink
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
900GB/s| GPU5["GPU 5"] GPU5 <--> GPU6["GPU 6"] GPU6 <--> GPU7["GPU 7"] GPU4 <--> GPU7 end Node1 <-->|NVSwitch / IB
跨节点互联| 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 进程都要调用）： ```python 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() ``` > [!note] NCCL 是什么 > **NCCL（NVIDIA Collective Communications Library）** 是 GPU 间集合通信的底层实现——它会探测实际硬件拓扑（NVLink / PCIe / InfiniBand），为每种 collective 选择 ring 或 tree 算法并生成通信 kernel。`torch.distributed` 只是它上面的一层薄封装；CPU 场景则退回 Gloo 后端（本地无 GPU 调试时很有用）。 **基本集合操作示例**： ```python 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**（讲义用代码验证上一讲的恒等式）： ```python 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`）： ```python 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 标称值对比。 > [!note] 带宽解读 > - **理论峰值**（H100 NVLink）：900 GB/s； > - **实测带宽**（all-reduce）：约 200–300 GB/s，且随张量增大而升高——小张量时固定延迟占主导，大张量才能吃满带宽； > - **差距原因**：ring/tree 算法需要多轮通信、协议开销、非直连跳数。做容量规划时永远用**实测**带宽，别用标称值。 > [!tip] 直觉 > 这一节的方法论与单卡 profiling 完全同构：先写 microbenchmark 校准硬件的真实能力（warmup → 同步 → 计时），再拿校准值去预估训练中的通信开销。后面三节对每种并行的「通信量分析」，都是用这里测出的 200–300 GB/s 代入算出来的。 --- ## 2. 数据并行：切分 batch 维度 ### 2.1 核心思想 **数据并行（Data Parallelism, DP/DDP）** 是最直观的并行策略： - **参数复制**：每个 GPU 持有模型的完整副本； - **数据切分**：把 batch 切成 $W$ 份，每个 GPU 处理一部分； - **梯度同步**：反向传播后，用 **all-reduce** 对所有 GPU 的梯度求平均，保证各副本走完全相同的优化轨迹。 ```mermaid graph LR subgraph "Rank 0" D0["数据
batch[0:32]"] --> M0["模型副本
θ"] M0 --> L0["loss₀"] L0 --> G0["∇₀"] end subgraph "Rank 1" D1["数据
batch[32:64]"] --> M1["模型副本
θ"] M1 --> L1["loss₁"] L1 --> G1["∇₁"] end subgraph "Rank 2" D2["数据
batch[64:96]"] --> M2["模型副本
θ"] M2 --> L2["loss₂"] L2 --> G2["∇₂"] end subgraph "Rank 3" D3["数据
batch[96:128]"] --> M3["模型副本
θ"] M3 --> L3["loss₃"] L3 --> G3["∇₃"] end G0 --> AR["all-reduce
求平均梯度"] G1 --> AR G2 --> AR G3 --> AR AR --> U["更新参数
θ ← θ - 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 的数据并行 ```python 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 ) ``` > [!warning] 易错点 > 梯度同步必须用 `ReduceOp.AVG`（或先 SUM 再除以 $W$）——如果只 SUM 不平均，等效学习率被放大 $W$ 倍，训练很可能发散。另外各 rank 的参数初始化必须一致（同随机种子，或由 rank 0 broadcast），否则「梯度相同」也无法保证「参数相同」。 ### 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 与反向传播**重叠**：某层梯度一算完就异步发起通信，藏在更早层的反向计算后面。 > [!tip] 直觉 > 数据并行的通信账本很简单：**每步付一次「模型大小」量级的固定通信税，与数据量无关**。所以它在「大 batch + 相对小模型」时几乎免费，在「小 batch + 大模型」时通信占比飙升。后面两种并行正是为了后一种场景准备的。 --- ## 3. 张量并行：切分宽度维度 ### 3.1 核心思想 **张量并行（Tensor Parallelism, TP）** 将模型的**每一层**按宽度切分——每个 GPU 只持有参数矩阵的一部分列。 **以 MLP 为例**（$y = \mathrm{GeLU}(xW)$，$W$ 形状 $[d, d]$）： ```mermaid graph LR subgraph "单 GPU（完整）" X1["x
[B, d]"] --> W1["W
[d, d]"] W1 --> Y1["y
[B, d]"] end subgraph "张量并行（4 GPU）" X2["x
[B, d]
（复制到所有 GPU）"] --> W0["W₀
[d, d/4]
GPU 0"] X2 --> W1p["W₁
[d, d/4]
GPU 1"] X2 --> W2["W₂
[d, d/4]
GPU 2"] X2 --> W3["W₃
[d, d/4]
GPU 3"] W0 --> Y0["y₀
[B, d/4]"] W1p --> Y1p["y₁
[B, d/4]"] W2 --> Y2["y₂
[B, d/4]"] W3 --> Y3["y₃
[B, d/4]"] Y0 --> AG["all-gather"] Y1p --> AG Y2 --> AG Y3 --> AG AG --> YF["y
[B, d]
（拼接）"] 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 的张量并行 ```python 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**，反向还要再来一轮——按激活算的通信量远超数据并行按梯度算的通信量； - 这就是张量并行「每层收税」的含义：通信频率是每层一次（而非每步一次），且不容易与计算重叠。 > [!note] Megatron-LM 的通信优化 > 朴素实现（如上）每个线性层后都 all-gather 一次。Megatron 的做法是**列切与行切配对**：第一个线性层按列切，GeLU 直接作用在列分片上（无需通信），第二个线性层按行切、吃进列分片、输出部分和——整个 MLP 块前向只在末尾做**一次 all-reduce**。再叠加序列并行时，这次 all-reduce 被等价替换为 all-gather + reduce-scatter（总字节数不变），顺便把 LayerNorm/dropout 的激活也切了（见 [[Lecture 7 · 并行 I：基础]] 第 7 节）。 > [!tip] 直觉 > 张量并行与数据并行是一对镜像：数据并行「复制参数、切分数据、同步梯度」，张量并行「切分参数、复制数据、同步激活」。选哪个取决于哪边的通信便宜——梯度大小固定为模型大小，激活大小正比于 batch；所以小 batch 大模型偏向张量并行，大 batch 小模型偏向数据并行。 --- ## 4. 流水线并行：切分深度维度 ### 4.1 核心思想 **流水线并行（Pipeline Parallelism, PP）** 将模型按**层**切分——每个 GPU 负责连续几层（一个 stage），激活在 GPU 间依次流动。 ```mermaid 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）**——朴素实现下游阶段等上游、上游算完等下游，大量时间空转： ```mermaid 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$ 个小批次流水线式填充，让各阶段错峰同时工作： ```mermaid 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 ``` > [!note] 气泡分析 > 设 $p$ 个流水线阶段、$m$ 个 micro-batch。注满流水线要 $p-1$ 个时隙、排空又要 $p-1$ 个，有效工作 $m$ 个时隙，所以气泡开销相对于有效计算为 > $$\text{bubble ratio} \approx \frac{p-1}{m}$$ > （占总时间的比例则为 $\frac{p-1}{m+p-1}$）。例如 $p=4$、$m=16$：相对开销 $3/16 \approx 19\%$，占总时间约 $16\%$。$m \gg p$ 是流水线高效的前提。 ### 4.2 最小实现：深度 MLP 的流水线并行 ```python 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 的额外计算换内存。 > [!tip] 直觉 > 流水线并行是「工厂流水线」的直译：单件产品（一个 batch）从头到尾串行加工快不起来，但把订单拆成许多小件（micro-batch）连续投料，所有工位就能同时忙碌。气泡对应「开工前的空转和收工前的清线」，订单越大（$m$ 越大）这部分固定成本摊得越薄。 --- ## 5. 3D 并行：组合三大策略 ### 5.1 为什么需要组合 **单一策略的瓶颈**： - **数据并行**：每步 all-reduce 梯度，通信量 $\propto$ 模型大小（70B 模型约 140 GB/步），且不切分任何内存； - **张量并行**：每层 all-gather/all-reduce 激活，通信量 $\propto$ 层数 × token 数 × hidden dim，出了 NVLink 就撑不住； - **流水线并行**：通信最省但有气泡，需要大 batch 摊平。 **组合策略（3D Parallelism）**：同时使用三种并行，让每种并行的通信落在它能承受的链路上： ```mermaid graph TB subgraph DP0["数据并行组 0"] subgraph PS00["流水线阶段 0"] TP00["张量并行
GPU 0-3"] end subgraph PS01["流水线阶段 1"] TP01["张量并行
GPU 4-7"] end end subgraph DP1["数据并行组 1"] subgraph PS10["流水线阶段 0"] TP10["张量并行
GPU 8-11"] end subgraph PS11["流水线阶段 1"] TP11["张量并行
GPU 12-15"] end end TP00 -.数据并行
梯度同步.-> TP10 TP01 -.数据并行
梯度同步.-> TP11 style TP00 fill:#c8e6c9,stroke:#2e7d32 style TP10 fill:#c8e6c9,stroke:#2e7d32 ``` **实测规模的参考配置**（Megatron 论文 [Narayanan et al. 2021](https://arxiv.org/abs/2104.04473) 中 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(激活) | 每微批次一次 | 跨节点，层数多 | > [!tip] 直觉 > 3D 并行的设计原则一句话：**通信频率与链路带宽逐级匹配**。最贵、最频繁的张量并行通信关进 NVLink（900 GB/s）；中频的流水线激活传递走跨节点网络；最外层数据并行每步只通信一次且可与反向重叠，对带宽最不挑剔。三个维度谁也替代不了谁，因为它们分别消解的是「单层放不下」「层数太多」「训练太慢」三种不同的约束。再算上 MoE 的专家并行（见 [[Lecture 4 · 混合专家模型 MoE]]）和长上下文的序列/上下文并行，实际系统是 4D 甚至 5D。 --- ## 6. 实战工具链 ### 6.1 PyTorch 原生支持 **DDP（DistributedDataParallel）**：数据并行的高级封装 ```python 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 实现 ```python 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 作为「声明式并行」的代表）。 --- ## 总结 ```mermaid mindmap root((并行 II
分布式训练实战)) 构建模块 集合通信原语 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$ 以最大化性能。 > [!note] 核心问题预热 > 花 100 万美元训练，是训 10B 模型 1T tokens，还是 1B 模型 10T tokens？缩放定律给出定量答案。 --- ## 复习自测 > [!question]- Q1：基准测试中 all-reduce 的「有效带宽」公式为什么有因子 $2(W-1)/W$？ > 环式 all-reduce 分 reduce-scatter 和 all-gather 两个阶段，每阶段中每个设备都要把 $\frac{W-1}{W}$ 份数据发出去（自己那份不用发），两阶段合计每卡收发 $2\frac{W-1}{W}s$ 字节。把这个「真实上链路的字节数」除以耗时，得到的带宽才能与 NVLink 标称值直接对比——这就是 nccl-tests 的 busbw 约定。单独测 reduce-scatter 时只有一个阶段，因子减半为 $\frac{W-1}{W}$。 > [!question]- Q2：算一算——70B 模型、BF16 梯度、8 卡 DDP、有效带宽 300 GB/s，每步梯度同步要多久？什么情况下这笔开销可以接受？ > 梯度大小 $G = 2 \times 70{\times}10^9 = 140$ GB；每卡收发 $2\frac{7}{8} \times 140 \approx 245$–$280$ GB，耗时约 $0.8$–$0.9$ 秒。若每步前向+反向计算耗时远大于此（大 batch），或者用 DDP 的分桶重叠把通信藏进反向传播，则可接受；小 batch 时通信占比过高，就该考虑增大本地 batch、梯度累积或改用模型并行。 > [!question]- Q3：概念辨析——数据并行同步「梯度」，张量并行同步「激活」，这个差别如何决定两者的适用场景？ > 梯度大小恒等于模型大小、与 batch 无关；激活大小正比于 token 数 × hidden dim × 层数、与模型深度和 batch 同时相关。因此大 batch 时数据并行的通信被计算摊薄，而张量并行的通信反而随 batch 线性膨胀且每层都在关键路径上。所以张量并行只在「模型单卡放不下」时开、且限制在 NVLink 节点内；数据并行则作为最外层默认维度。 > [!question]- Q4：为什么流水线并行的最小实现用 send/recv 而不用 collective？micro-batch 数从 4 增到 16，气泡怎么变？还有什么代价？ > 流水线的通信天然是「阶段 $i$ → 阶段 $i{+}1$」的一对一传递，不涉及全体设备，点对点原语正好匹配。气泡相对开销从 $(p-1)/4$ 降到 $(p-1)/16$（4 阶段时从 75% 降到约 19%）。代价：micro-batch 变小可能拖低单卡 kernel 效率；朴素调度下需同时保存的在途激活份数增加（1F1B 可将其压到 $p$ 份量级）；通信次数线性增加。 > [!question]- Q5：设计题——两台 8×H100 节点（共 16 卡）训练一个单卡放不下的稠密模型，如何分配 TP/PP/DP？说明理由。 > 参考方案：TP=8 放节点内（每层的高频 all-reduce 走 NVLink），PP=2 跨节点（每 micro-batch 只传一次层间激活，耐受节点间低带宽），DP=1（卡数已用尽）。若模型经 TP=8 切分后每卡显存仍富余，可改为 TP=4 × PP=2 × DP=2，用数据并行提高吞吐；此时 DP 的梯度 all-reduce 也应尽量约束在节点内或与反向重叠。核心逻辑始终是：通信频率最高的并行维度占据带宽最高的链路。 --- ## 参考资料 **论文**： - **Megatron-LM**（NVIDIA, 2019）：[Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://arxiv.org/abs/1909.08053) - 提出张量并行（Tensor Parallelism）在 Transformer 中的实现 - **GPipe**（Google, 2019）：[GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism](https://arxiv.org/abs/1811.06965) - 流水线并行 + 微批次 + 梯度累积 - **ZeRO**（Microsoft, 2020）：[ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://arxiv.org/abs/1910.02054) - 数据并行的内存优化（分片优化器状态、梯度、参数） - **PyTorch FSDP**（Meta, 2021）：[Fully Sharded Data Parallel: faster AI training with fewer GPUs](https://engineering.fb.com/2021/07/15/open-source/fsdp/) - ZeRO-3 的 PyTorch 原生实现 - **3D 并行实测**（NVIDIA/Stanford/MSR, 2021）：[Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM](https://arxiv.org/abs/2104.04473) - 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 流调度，是从「能跑」到「跑满」的关键一步 --- > [!tip] Lecture 8 核心收获 > 上一讲画了作战地图，这一讲拿起工兵铲挖战壕——用不到 200 行 PyTorch 代码，把数据并行 / 张量并行 / 流水线并行从图纸变成可执行的 Python 函数。下一讲，我们将用缩放定律回答「训练多大的模型最划算」。