Lec20 并行与扩展

本讲解决 LLM 预训练的核心系统问题：如何高效利用多 GPU 来训练更大的模型、处理更多的数据。关键问题在于如何在设备间分配工作，同时最小化通信和空闲气泡造成的浪费。

动机

预训练同时受益于更多数据和更大模型——两者都能降低训练损失。问题在于扩展需要高效利用多 GPU，这引入了三个相互制约的因素：

内存占用：所有训练状态（权重、梯度、优化器状态、激活值）必须装入 GPU 显存
计算效率：GPU 应把大部分时间花在计算上，而非等待
通信开销：GPU 间的同步会导致空闲

正确配置的影响是巨大的。课件展示了同一个 3.57B 模型在 32 节点上的两种配置：一种达到 17.4% MFU、1.7M tok/s，另一种仅 0.7% MFU、60K tok/s——仅因并行策略不同，吞吐量相差 28 倍。

单 GPU 训练基础

计算量

训练计算量以浮点运算（FLOPs）衡量。一次前向和反向传播的计算量为：

\[ \text{FLOPs} = 6 \times \text{model\_parameters} \times \text{token\_batch\_size} \]

系数 6 来自：前向传播中乘加运算 2x，反向传播 4x（两次反向计算各 2x）。

模型浮点运算利用率（MFU）衡量硬件的使用效率：

\[ \text{MFU} = \frac{\text{实际 FLOPS}}{\text{理论峰值 FLOPS}} \]

作为参考，H100 SXM 在 BFloat16 Tensor Core 下的理论峰值为 1,979 teraFLOPS。实际训练通常达到 30-45% MFU，原因在于通信开销、内存带宽限制和空闲时间。

内存占用

一个训练步骤需要存储四类张量：

\[ \text{峰值内存} = \text{model\_bf16} + \text{model\_fp32} + \text{grads\_fp32} + \text{optim\_states} + \text{activations} \]

使用混合精度 Adam 时，每个参数的大致开销：

组件	每参数字节数
BF16 模型权重	2
FP32 主权重	4
FP32 梯度	4
Adam 动量（FP32）	4
Adam 方差（FP32）	4
总计（不含激活值）	18

对于 7B 模型，仅参数/梯度/优化器就需要约 126 GB——已经超过单张 H100 的 80 GB。70B 模型需要约 1,120-1,400 GB。

激活值内存

激活值内存与 batch size 线性相关，与序列长度二次相关：

\[ m_{\text{act}} = L \cdot seq \cdot bs \cdot h \cdot \left(34 + \frac{5 \cdot n_{\text{heads}} \cdot seq}{h}\right) \]

其中 \(L\) 为层数，\(seq\) 为序列长度，\(bs\) 为 batch size，\(h\) 为隐藏维度，\(n_{\text{heads}}\) 为注意力头数。二次项来自注意力分数矩阵（\(seq \times seq\)）。在长序列（8K+）下，激活值主导内存占用。

激活值重计算

激活值重计算（梯度检查点）不在前向传播中保存所有激活值，而是仅保存某些激活值作为”检查点”，在反向传播中重新计算其余部分。

权衡很明确：用更多计算换更少内存。选择性重计算针对最耗内存的激活值（如注意力分数），同时保留较便宜的激活值，在两个极端之间取得平衡。

梯度累积

梯度累积将大 batch 拆分为小的 micro-batch，依次执行前向/反向传播并在更新前平均梯度：

\[ bs = gbs = mbs \cdot grad\_acc \]

这使你可以在恒定内存下模拟大 batch size。内存消耗由 micro-batch size 决定，而有效 batch size 等于 micro-batch size 乘以累积步数。预训练通常使用每 batch 400-6000 万 token。

数据并行

基本思路

数据并行（DP）在每个 GPU 上复制完整模型。每个 GPU 处理不同的 micro-batch，独立计算梯度，然后所有 GPU 通过 all-reduce 操作平均梯度后更新权重。

全局 batch size 变为：

\[ \text{global batch size} = mbs \cdot grad\_acc \cdot dp \]

其中 \(dp\) 为数据并行副本数。

通信：集合操作

多 GPU 训练依赖集合通信原语：

操作	模式
Send/Recv	两个 rank 之间的点对点传输
Scatter	一个 rank 向所有 rank 分发数据块
Gather	一个 rank 从所有 rank 收集数据块
Broadcast	一个 rank 向所有 rank 发送相同数据
Reduce	所有 rank 贡献数据，一个 rank 获得聚合结果
All-Reduce	所有 rank 贡献数据，所有 rank 获得聚合结果
All-Gather	所有 rank 贡献数据块，所有 rank 获得完整数据

朴素数据并行中，整个反向传播完成后发起一次 all-reduce。这意味着 GPU 在整个通信阶段处于空闲状态。

重叠 + 分桶

关键优化是将通信与计算重叠：

梯度就绪后立即启动 all-reduce — 不等待完整的反向传播
将梯度分组到桶中 — 每个桶发起一次 all-reduce 而非每个参数一次

由于反向传播从最后一层开始计算梯度，后面层的 all-reduce 可以与前面层的梯度计算重叠，将大部分通信延迟隐藏在有用的计算背后。

扩展行为

数据并行不会减少每 GPU 内存——每个 GPU 仍持有模型、梯度和优化器状态的完整副本。它仅通过并行处理更多数据来提升吞吐量。

随着 GPU 数量增加，由于通信开销增长，每 GPU 吞吐量下降。课件展示 3B 模型从 DP=8 时的约 38K tokens/sec/GPU 降至 DP=256 时的约 22K tokens/sec/GPU（-40.6%）。

实例计算

假设全局 batch size 为 400 万 token，序列长度 4,000： - 序列数：4M / 4K = 1,024 - 若单 GPU 可容纳 mbs=2 个序列，128 个 GPU：2 × 128 = 256 个序列/步 - 需要 grad_acc = 1024 / 256 = 4 个累积步

若改为 512 个 GPU：2 × 512 = 1,024 个序列/步，grad_acc = 1（无需累积，每步更快）。

张量并行

基本思路

当模型太大无法放入单个 GPU 时，张量并行（TP）将单个权重矩阵拆分到多个 GPU 上。每个 GPU 计算矩阵乘法的一部分，然后通过通信合并结果。

列并行

将权重矩阵 \(W\) 按列拆分。每个 GPU 持有一个列块，独立计算 \(X \cdot W_i\)。输入 \(X\) 广播到所有 GPU，部分输出 \(Y_i\) 通过 all-gather 重建完整输出。

2 GPU 拆分 \(Y = XW\)： - GPU 0 计算：\(Y_0 = X \cdot W_0\)（前半列） - GPU 1 计算：\(Y_1 = X \cdot W_1\)（后半列） - All-gather 得到：\(Y = [Y_0, Y_1]\)

行并行

将权重矩阵按行拆分（相应地将输入按列拆分）。每个 GPU 计算部分结果，然后通过 all-reduce（求和）合并。

2 GPU 拆分： - GPU 0 计算：\(Y_0 = X_0 \cdot W_0\) - GPU 1 计算：\(Y_1 = X_1 \cdot W_1\) - All-reduce 得到：\(Y = Y_0 + Y_1\)

应用于 Transformer 层

前馈层第一个线性层（\(A\)）使用列并行，第二个线性层（\(B\)）使用行并行。这避免了两者之间的中间 all-reduce/all-gather——列并行的输出自然成为行并行的输入：

\[ Y = \text{GeLU}(XA), \quad Z = \text{Dropout}(YB) \]

\(A = [A_1, A_2]\)（列拆分）、\(B = \begin{bmatrix} B_1 \\ B_2 \end{bmatrix}\)（行拆分），每个 GPU 计算 \(\text{GeLU}(X \cdot A_i) \cdot B_i\)，最后只需一次 all-reduce。

注意力层按注意力头自然拆分。Q、K、V 投影按列拆分，每个 GPU 独立处理其子集的注意力头，输出投影按行拆分。

权衡

张量并行减少了权重、梯度、优化器状态和激活值的每 GPU 内存。70B 模型在无并行时需要 140 GB，TP=8 或 TP=16 可降至 80 GB 以内。

代价是通信。每个 Transformer 层在前向和反向传播中都需要 all-reduce 或 all-gather 操作。吞吐量随 TP 增大急剧下降：3B 模型 TP=2 时约 13K tokens/sec/GPU，TP=32 时降至约 4.5K（-65.6%）。

跨节点通信带宽大幅下降（节点内 436 GB/s vs 64 节点时约 34 GB/s），因此张量并行应限制在单节点内（通常每节点 8 个 GPU）。

流水线并行

基本思路

流水线并行（PP）按层将模型拆分到多个 GPU 上。GPU 1 持有第 1-4 层，GPU 2 持有第 5-8 层，以此类推。数据依次流过流水线。

这按比例减少每 GPU 内存——8B 模型在 PP=8 下从约 140 GB 降至约 40 GB/GPU。

气泡问题

朴素方法有巨大的低效：当 GPU 1 处理前向传播时，GPU 2-4 处于空闲；当 GPU 4 处理反向传播时，GPU 1-3 空闲。这些浪费的时间称为流水线气泡。

One-Forward-One-Backward（1F1B）调度

1F1B 调度通过在 micro-batch 间交替前向和反向传播来减少气泡。每个 GPU 不等所有前向传播完成就开始反向传播，而是在最后一个流水线阶段完成后立即开始。

更多 micro-batch 使气泡占总时间的比例更小。课件显示 PP=2 时吞吐量约 15K tokens/sec/GPU，PP=32 且少量 microbatch 时下降 -44.3%。32 个 microbatch 时，PP=32 的退化为 -52.4%，仍然显著。

流水线并行最适合跨节点使用，因为它仅需点对点通信（向下一阶段发送激活值、从下一阶段接收梯度），不像张量并行需要 all-reduce。

内存优化：ZeRO

冗余问题

标准数据并行中，每个 GPU 存储模型参数、梯度和优化器状态的完整副本。对于 7B 模型使用 Adam，这意味着约 126 GB 在每个 GPU 上重复存储——极其浪费。

ZeRO 阶段

零冗余优化器（ZeRO）渐进式地将训练状态分片到各 GPU：

阶段	分片内容	内存节省	通信开销
基线（DP）	无	无	低（仅 all-reduce）
ZeRO-1	优化器状态	优化器内存约 4x 减少	低
ZeRO-2	+ 梯度	进一步减少	中等
ZeRO-3	+ 参数	最大减少	高（每层 all-gather）

ZeRO-3 工作原理

每个 GPU 仅存储 \(1/N\) 的参数（\(N\) 为 GPU 数量）。在前向/反向传播的每一层： 1. GPU 发起 all-gather 从所有对端获取该层的完整参数 2. 使用完整参数计算激活值/梯度 3. 释放获取的参数，继续下一层

与张量/流水线并行的关键区别：ZeRO 分片的是内存，而非计算。每个 GPU 仍在完整数据上执行完整计算——只是不同时存储所有内容。这意味着 ZeRO-3 通信开销最高（前向和反向传播中每层都需要 all-gather），但提供最大的内存节省。

策略选择

方法总结

策略	核心思想	权衡	最佳场景
数据并行（DP）	在 batch 维度并行	冗余内存；模型须能放入 GPU	能放入 GPU 内存的标准模型
张量并行（TP）	在隐藏维度并行	高通信（每层 all-reduce）	大型层；节点内并行
流水线并行（PP）	在模型维度并行	流水线气泡浪费时间	大型深层模型；跨节点并行
ZeRO	在 DP 中分片模型/优化器/梯度	高通信（all-gather）	无法放入 GPU 内存的大模型

决策框架

配置训练时的目标： 1. 将模型装入内存 — 使用 TP、PP 或 ZeRO 减少每 GPU 内存 2. 满足目标全局 batch size — 使用 DP 和梯度累积 3. 最大化训练吞吐量 — 最小化通信开销和空闲时间

实践中，这些策略需要组合使用。典型配置可能为： - 节点内使用 TP（如 TP=4 或 TP=8），利用快速节点内通信减少内存 - 跨节点使用 PP（如 PP=2 或 PP=4），仅需点对点跨节点通信 - 剩余 GPU 使用 DP 扩展吞吐量 - ZeRO-1 分片优化器状态，通信开销不大

最佳配置实验

Ultra-Scale Playbook 的基准测试展示了最优配置如何随模型大小和集群规模变化（1M token GBS，序列长度 4096，H100 节点）：

小模型（1.3B）少量节点：纯 DP 效果好（如 DP=4, TP=1, PP=2, MFU ~38%）
中等模型（3.5-8.8B）：节点内 TP=4 + 适度 PP + DP 可达 ~40-46% MFU
大模型（80B）：需要激进的 TP（TP=4-16）+ PP（PP=4-16）+ ZeRO-1，MFU 根据集群规模在 ~15-45% 之间
更多节点通常降低 MFU：1.3B 模型从 1 节点到 64 节点，MFU 从约 45% 降至约 5%，通信开销占主导

训练框架

两个主要框架实现了这些并行策略：

torchtitan：PyTorch 原生平台，支持 FSDP2、Tensor Parallel、Pipeline Parallel、Context Parallel、激活检查点、Float8、梯度累积等
Megatron-LM：NVIDIA 的 GPU 优化库，Megatron Core 提供内核、张量/流水线并行、分布式训练（FSDP、DDP）和后训练支持

Lec20 要点总结

内存是第一个瓶颈：7B 模型仅训练状态就需要约 126 GB；激活值重计算和梯度累积有帮助但对大模型无法根本解决
数据并行扩展吞吐量，不减内存：每个 GPU 仍持有完整模型；通信开销随 GPU 数量增长
张量并行以通信换内存：在节点内拆分权重矩阵；应限制在快速节点内链路中
流水线并行引入气泡：跨 GPU 拆分层产生空闲时间；1F1B 调度和更多 microbatch 减少但不能消除气泡
ZeRO 消除冗余：渐进式分片优化器状态、梯度和参数，以通信换内存节省；与 TP/PP 本质不同，分片的是内存而非计算

总结

主题	核心思想
计算量与 MFU	每步 6 × 参数 × token 数；MFU = 实际/峰值 FLOPS；实际训练约 30-45%
内存分解	混合精度 Adam 每参数 18 字节；激活值随序列长度二次增长
激活值重计算	用约 33% 更多计算换取大幅减少的激活值内存
梯度累积	通过 micro-batch 平均在恒定内存下模拟大 batch
数据并行	复制模型，拆分数据；all-reduce 梯度；通过分桶重叠通信
张量并行	拆分权重矩阵（列/行）；因通信开销应限于节点内
流水线并行	跨 GPU 拆分层；1F1B 调度减少气泡；适合跨节点
ZeRO	跨 DP rank 分片优化器/梯度/参数；内存分片而非计算分片
策略选择	组合 TP（节点内）+ PP（跨节点）+ DP（其余）+ ZeRO-1；优化 MFU

核心要点：训练大型 LLM 需要组合多种并行策略——每种解决问题的不同维度（数据、隐藏维度、层、内存），最优组合取决于模型大小、集群拓扑和 batch size 约束。

参考文献

The Ultra-Scale Playbook: Training LLMs on GPU Clusters (HuggingFace)
torchtitan: github.com/pytorch/torchtitan
Megatron-LM: github.com/NVIDIA/Megatron-LM
Rajbhandari et al., “ZeRO: Memory Optimizations Toward Training Trillion Parameter Models,” 2020
Shoeybi et al., “Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism,” 2019

本文基于 CMU 11-711 Advanced NLP 课程资料，由 Sean Welleck 讲授（Lecture 20: Parallelism and Scaling）。