1. Outline and Goals从优化单个 GPU 的吞吐量到理解训练超大模型所需的复杂性和细节。2. Basics of Networking for LLMs2.1 Hardware无论是从算力还是内存的角度考虑single GPU 都无法满足训练需求。这个图类似 Lecture 1 中1.2.2.2 Parallelism可直接参考 CPU 经过 PCLe Switch 中转通过 PCLe 与 GPU 相连GPU 之间经过 NV Switch 中转通过 NV Link 相连。PLX 可以理解为一个 PCle SwitchHCA 负责将服务器连接到外部的 InfiniBand 高速集群网络实现多机通信构建一个大规模计算集群。2.2 Collective Communication Operations一组进程或GPU之间进行一对多、多对一或多对多的数据交换和协同计算Root根:在操作中拥有主数据或负责汇总的特殊进程。Reduce规约将多个数据通过某种操作合并成一个数据。All reduce原理上为 Reduce 和 Broadcast 的组合对每个进程进行规约每个进程得到相同的最终结果。计算成本约为待规约数据的两倍。模型训练过程中反向传播中的梯度同步就是 All reduce。但按 Reduce Broadcast 方式计算性能太差—— Root 进程需要同时接受所有进程的数据还需要同时发送结果给所有进程。流量都集中在 Root 进程多块网卡的带宽未充分利用。工程实现上All reduce Reduce Scatter All Gather将通信荷载平均分配给所有节点而不是让 Root 成为性能瓶颈。通信成本仍为 2Ψ2.3 TPUs vs GPUsGPU 实现多台机器网络互联单台机器通常配备 8 块 GPU这些节点通过 Switch 高速互联。机器之间通过叶交互机相连通常最多可扩展至 256 块 GPU 规模即最多有32台机器由叶交互机彼此相连。若要再拓展到更大规模的计算集群叶交互机之间以脊交互机相连。之所以 Reduce Brandwidth 会下降是因为 8 块 GPU 时所有通信走的是 NV Link速度极快而跨机器通信只能走 InfiniBand 网络速度远远下降。TPU 实现网络互联更改了机器之间的组网方式每块芯片都能与相邻芯片实现高速互联具有极高的可扩展性称为“环状网络Toroidal Mesh”。但问题是每块芯片只能与相邻节点直接通信。2.4 Recap3. Different Forms of Parallel LLM Training3.1 Data Parallelism每张卡都拥有完整的模型副本但每张卡只负责处理一部分数据。3.1.1 Naive Data Parallelism数据并行的基础是随机梯度下降算法Stochastic Gradient DescentSGD。计算扩展性compute scale: Batch Size 为 B共有 M 张 GPU则每张 GPU 负责 B / M 个样本。计算负载被平均分配给了所有 GPU计算能力随 GPU 数量近似线性扩展。通信开销communication overhead每 Batch 传递 2 倍数量的参数来自 All Reduce 进行梯度同步。由于通信开销只与模型参数量成正比为一个固定值而计算开销与 Batch Size 成正比当 Batch 足够大时计算开销远大于通信开销设定合理。内存扩展性memory scaling差每张 GPU 都需要存储完整的模型参数。如下图所示主权重master weights一份以更高精度存储的模型参数。前向传播时复制一份低精度的模型参数进行计算 pred 结果反向传播时计算梯度都只用低精度参数。只有最后优化器更新才会修改主权重参数。在混合精度训练过程中主要占显存的不是模型参数本身而是优化器状态Optimizer State。在典型的 Adam 混合精度训练过程中平均每个模型参数需要 16 Bytes 的显存Lecture 2 3.5写的是 AdaGrad 优化器3.1.2 ZeRo在参数内存方面优化器状态占用的内存为主要。Ψ模型参数数量 N_dGPU 数量在实际操作中并非需要再每台 GPU 上保存一份优化器状态以至于梯度和模型参数。基于这种思想发展出了一套分布式训练的内存优化算法ZeRO (Zero Redundancy Optimizer)通过切分Sharding来消除数据并行中模型状态的冗余用通信换显存。①. ZaEO-1 : Optimizer State Sharding将优化器状态主权重 动量 m 动量 v切分到不同的 GPU 上而非每张卡存放一份完整的。假设现有 4 张 GPU 以及模型具有 12 个参数。step 0初始阶段——所有 GPU 都具有完整的模型参数12个和梯度占位符12个step 1前向传播——根据自己的模型参数12个计算预测值step 2反向传播——根据预测值以及真实值计算梯度step 3梯度同步——传统标准分布式数据并行Distributed Data Parallelism做法All Reduce将所有 GPU 的 12 个梯度求平均每张 GPU 获得完整的 12 个平均梯度。ZeRO-1 做法Reduce Scatter不将结果全部发给所有 GPU只将每张 GPU 负责的那部分梯度规约给对应的 GPU。例如 GPU 0 就只负责参数 0、1、2。step 4优化器更新——每张 GPU 用上一步得到的对应参数的平均梯度以及本地存储的优化器状态主权重动量 m动量 v更新对应的参数。例如 GPU 0 完成对参数 0、1、2 的更新。其他参数不变。step 5参数同步—— All Gather每张 GPU 分别将更新后的模型参数广播给其他参数。这一步完成后所有 GPU 上的 12 个参数都同步成了最新值。如上文所说All Reduce 的通信成本 Reduce Scatter All Gather。所以 ZeRO-1 的 Reduce Scatter 和 All Gather 的组合操作并未引入新的通信成本。相当于免费获得了更低的内存占用。②. ZeRO-2 Gradient Sharding在 ZeRO-1 的基础上将 gradient 也切分掉。step 0初始阶段每张 GPU 只具有完整的模型参数12个以及 3 个梯度占位符step 1前向传播同上step 2反向传播同上——各 GPU 计算出完整的 12 个梯度step 3梯度同步——第一步同 ZeRO-1利用 Reduce Scatter 为每张 GPU 计算出对应部分的梯度第二步对于每张 GPU 立刻释放无关梯度的显存。释放梯度会产生一些额外开销但可以忽略step 4优化器更新同上step 5参数同步同上③. ZeRO-3 Parameter Sharding在 ZeRO-2 的基础上把模型参数也切分了。在Pytorch 上的工程实现FSDPFully Sharded Data Parallelstep 0初始阶段每张 GPU 只具有各自负责的部分模型参数3个以及梯度占位符3个step 1前向传播需要完整的参数但本地没有——每次计算前利用 All Gather 先收集过来计算完毕后再释放掉。相较于 ZeRO-1 和 ZeRO-2 这里引入了新的通信开销step 2反向传播需要完整的参数但本地没有——同样是一次 All Gatherstep 3梯度同步同上一次 Reduce Scatterstep 4优化器更新同上step 5参数同步——可以理解成没有这一步了因为模型训练是个循环step 4 后直接 step 1在 step 1 处进行了 All Gather。所以 ZeRO-3 的通信消耗为 3Ψ两次 All Gather 以及一次 Reduce Scatter。以 Ψ 为通信代价换来了原内存占用 / N 的内存收益。由于通信时间与计算时间可以重叠所以 ZeRO-3 的实际运行开销极低。理想状态是当 GPU 在运算的同时后台通信就像预加载机制般持续工作。需要某块数据时它早已传输到位随时可供调用。以一个计算为例在 ZeRO-3 / FSDP 中前向传播和反向传播的通信开销几乎可以被计算完全“掩盖”掉实际额外耗时极低3.1.3 Problem of Data Parallelism①.数据并行的天花板是 Batch Size 大小因为并行度无法超过 Batch Size 规模本身。但在训练过程中无法一味增大 Batch Size当 Batch Size 超过某个临界值后优化效率的提升幅度会迅速衰减。所以单靠数据并行无法实现无限扩展的并行训练能力。②. 数据并行没有解决激活内存的问题即使是 ZeRO-3后众 GPU 也在 All Gather 参数后每张 GPU 都具备完整参数计算出完整的激活值保存在 GPU 上。一次训练过程中内存占据情况见 3.3 中图3.2 Model Parallelism把一个模型切分成多个部分分布到不同的 GPU 上每张 GPU 只负责模型的一部分计算。3.2.1 Pipeline Parallelism将模型的不同层分配给不同的 GPU。前向传播中 GPU 之间通过传递激活值实现协同计算反向传播时 GPU 传递梯度值。问题GPU 闲置严重需要等前一步 GPU 计算完毕n 个 GPU每个 GPU 的激活时间只有 1/n解决方法Micro-batches系统开销与有效计算的比值(m - 1) / n m 为 Pipeline 个数n 为 micro-batch 个数当 Batch Size 足够大时Data Parallelism 和 Pipeline Parallelism 都不失为一种好的并行方法。将两者组合以获得最好的优化空间。但同理由于 Batch Size 有限无法扩展到任意规模的并行。Pipeline Parallelism 是不是本身就带有 Data Parallelism如果有 8 张 GPU但 PP size 2即会分出 4 组每组 2 张 GPU 分别进行一次完整的所有层的计算。那分出的 4 组不就天然 DP 了吗Pipeline Parallelism 的优势1. 相较于 Data ParallelismPipeline Parallelism 能显著降低内存开销。两者都切分了参数虽然切法不一样PP 进一步切分了激活值2. PP 的通信效率高因为 GPU 之间仅传递激活值和梯度。PP 常用于网络中的低速链路场景例如节点间通信、跨机架甚至跨数据中心的场景。正如上文提到的当计算集群超过 256 块 GPU 时需要脊交互机连接时此时就适合 PP 应用。Zero_Bubble Pipeline反向传播中的操作可以分为两个阶段B 阶段Backward for Input计算损失函数对当前层输入的梯度。这个结果需要传递给前一层用于计算该层的梯度。W 阶段Backward for weights计算损失函数对当前层参数的梯度。这个结果用于本地更新模型权重。该步骤只需要在 B 阶段后无其他的时序依赖故可以安排到 Bubble 中。1F1B 如果所示系统开销与有效计算比值为 3/8ZB-H1 没有额外增加激活内存气泡进一步减少图中为 1/8ZH-H2 牺牲 2 倍激活内存完全消除气泡但如此做的复杂度极高平常调用 loss.backward() 会自动计算梯度而 Zero Bubble 需要某个 GPU 在特定时刻计算某个 micro-batch 的 B 或 W 操作。为此系统必须精确控制每个 micro-batch 的数据流向并主动打断 PyTorch 默认的执行流程在计算图完成一次操作后显式地触发下一次调度。3.2.2 Tensor Parallelism把单层的权重矩阵切成多块分布到多张 GPU 上共同完成这一层的计算。有点像 Lecture 5 中 3.6 Tiling只不过 Tiling 将大矩阵切成小片以便放进 GPU 的共享内存Shared Memory中减少对全局显存的访问提升单卡计算效率数据始终在一张卡上而 Tensor Parallelism 数据被物理上移到了不同的卡上通过通信来交换中间结果前向传播和反向传播都需要进行一次 All Reduce。而且上图中的 f 和 g 分别是一个同步屏障确保所有参与计算的 GPU 都完成了某个关键步骤后再允许任何一张 GPU 继续执行后续操作。代价强制快的 GPU 等待慢的 GPU导致计算资源的闲置。因此张量并行的通信频率极高每层都要同步因此通常只在单机内部使用由 NVLink 高速互联。3.3 Activation Parallelism一次训练过程中内存占据随时间变化图3.3.1 no parallelismActivationMemoryPerTransformerLayer :​​​​​​​ ​​​​​​​ ​​​​​​​ ​​​​​​​(a: number of attention heads; s: sequence length; h: hidden dimension size)第一项来自于 MLP 等逐元素操作point-wise operations后一项来自于 softmax 和其他二次项计算。3.3.2 Tensor Parallel​​​​​​​ ​​​​​​​ ​​​​​​​ ​​​​​​​( t : tensor parallel size)其中第一项与 t 无关其组成部分大致为 LayerNorm、Dropout、注意力机制的输入以及 MLP 等。这些操作都会随着模型规模增长而膨胀而且它们的并行化效果并不理想。所以对这些项进行操作就是处理的核心。3.3.3 Sequence ParallelTensor parallel 切的是权重矩阵的列hidden layer每张 GPU 仍知道所有的 Token但只有一半特征。所以像 LayerNorm 这种需要跨所有列计算的操作无法在单卡上独立完成。Tensor parallel 无法对其优化。Sequence parallel 切的是行。每张 GPU 只拿到一部分 Token但每个 Token 都具有完整特征。在前向传播中g 为 All Gatherg_hat 为 Reduce Scatter而在反向传播中刚好相反。activation recomputation 参考 Lecture 5 中 3.43.4 Other Parallelism3.4.1 Context Parallel / Ring AttentionQ、K、V 都是 [s, h]大小的矩阵s 是序列长度h 是隐藏维度。当 s 很大时Q K^T 会生成一个 [s, s] 的巨大矩阵单卡显存放不下。Context Parallel让每张 GPU 只存自己那一段的 Q但 K 和 V 在环形网络上轮流传递给每一张 GPU。3.4.2 Expert ParallelLecture 4 中 2.3.4将一个大的 MLP 变成多个每次稀疏激活若干个的小 MLP专家4. Scaling and Training Big LMs with Parallelism在并行策略中Batch Size 为关键参数根据 Batch Size 与 GPU 数量之比的不同最佳并行策略也会相应变化。当 Batch Size 远低于 GPU 数量时效率必然低下——计算时间短而通信时间为固定的系统受限于通信瓶颈图中黑色区。通信时间无论 Batch Size 是 1 或 1024每次迭代都需要同步梯度。All-Reduce 通信的数据量为模型参数量 Ψ 的 2 倍与 Batch Size 完全无关。而计算时间随 Batch Size 线性变化而随着 Batch Size 增加可采用的混合并行策略才有意义计算耗时超过通信耗时。将这些并行技术结合起来就得到了业界所称的 3D 或 4D 并行方案。1. 前提保证模型和参数值能装入内存。优先用 Model ParallelTensor Parallel单机之间以及 Pipeline Parallel机器之间2. 模型装得下后用 Data Parallem 进行拓展。3. 如果 Batch Size 太小此时通信开销会占主导。在本地算多个 batch 的梯度累加后再统一进行一次 All-Reduce 同步。相当于用计算时间换取通信频率的降低让通信效率变高。