企业官网建设流程全解析

1. Outline and Goals

从优化单个 GPU 的吞吐量到理解训练超大模型所需的复杂性和细节。

2. Basics of Networking for LLMs

2.1 Hardware

无论是从算力还是内存的角度考虑，single GPU 都无法满足训练需求。

（这个图类似 Lecture 1 中1.2.2.2 Parallelism，可直接参考 ）

CPU 经过 PCLe Switch 中转，通过 PCLe 与 GPU 相连；GPU 之间经过 NV Switch 中转，通过 NV Link 相连。

PLX 可以理解为一个 PCle Switch；HCA 负责将服务器连接到外部的 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 GPUs

GPU 实现多台机器网络互联：

单台机器通常配备 8 块 GPU，这些节点通过 Switch 高速互联。机器之间通过叶交互机相连，通常最多可扩展至 256 块 GPU 规模（即最多有32台机器由叶交互机彼此相连）。若要再拓展到更大规模的计算集群，叶交互机之间以脊交互机相连。

之所以 Reduce Brandwidth 会下降，是因为 8 块 GPU 时所有通信走的是 NV Link，速度极快；而跨机器通信只能走 InfiniBand 网络，速度远远下降。

TPU 实现网络互联：

更改了机器之间的组网方式，每块芯片都能与相邻芯片实现高速互联，具有极高的可扩展性，称为“环状网络（Toroidal Mesh）”。但问题是每块芯片只能与相邻节点直接通信。

2.4 Recap

3. Different Forms of Parallel LLM Training

3.1 Data Parallelism

每张卡都拥有完整的模型副本，但每张卡只负责处理一部分数据。

3.1.1 Naive Data Parallelism

数据并行的基础是随机梯度下降算法（Stochastic Gradient Descent，SGD）。

计算扩展性（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_d：GPU 数量）

在实际操作中，并非需要再每台 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 上的工程实现：FSDP（Fully Sharded Data Parallel）

step 0：初始阶段，每张 GPU 只具有各自负责的部分模型参数（3个）以及梯度占位符（3个）

step 1：前向传播，需要完整的参数，但本地没有——每次计算前利用 All Gather 先收集过来，计算完毕后再释放掉。相较于 ZeRO-1 和 ZeRO-2 这里引入了新的通信开销

step 2：反向传播，需要完整的参数，但本地没有——同样是一次 All Gather

step 3：梯度同步，同上一次 Reduce Scatter

step 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 Parallelism，Pipeline Parallelism 能显著降低内存开销。两者都切分了参数（虽然切法不一样），PP 进一步切分了激活值

2. PP 的通信效率高，因为 GPU 之间仅传递激活值和梯度。PP 常用于网络中的低速链路场景，例如节点间通信、跨机架甚至跨数据中心的场景。

正如上文提到的，当计算集群超过 256 块 GPU 时，需要脊交互机连接时，此时就适合 PP 应用。

Zero_Bubble Pipeline

反向传播中的操作可以分为两个阶段：

B 阶段（Backward for Input）：计算损失函数对当前层输入的梯度。这个结果需要传递给前一层，用于计算该层的梯度。

W 阶段（Backward for weights）：计算损失函数对当前层参数的梯度。这个结果用于本地更新模型权重。该步骤只需要在 B 阶段后，无其他的时序依赖，故可以安排到 Bubble 中。

1F1B 如果所示，系统开销与有效计算比值为 3/8

ZB-H1 没有额外增加激活内存，气泡进一步减少（图中为 1/8）

ZH-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 parallelism

ActivationMemoryPerTransformerLayer :

​​​​​​​ ​​​​​​​ ​​​​​​​ ​​​​​​​

(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 Parallel

Tensor parallel 切的是权重矩阵的列（hidden layer），每张 GPU 仍知道所有的 Token，但只有一半特征。所以像 LayerNorm 这种需要跨所有列计算的操作，无法在单卡上独立完成。Tensor parallel 无法对其优化。

Sequence parallel 切的是行。每张 GPU 只拿到一部分 Token，但每个 Token 都具有完整特征。

在前向传播中，g 为 All Gather，g_hat 为 Reduce Scatter；而在反向传播中刚好相反。

activation recomputation 参考 Lecture 5 中 3.4

3.4 Other Parallelism

3.4.1 Context Parallel / Ring Attention

Q、K、V 都是 [s, h]大小的矩阵（s 是序列长度，h 是隐藏维度）。当 s 很大时，Q @ K^T 会生成一个 [s, s] 的巨大矩阵，单卡显存放不下。

Context Parallel：让每张 GPU 只存自己那一段的 Q，但 K 和 V 在环形网络上轮流传递给每一张 GPU。

3.4.2 Expert Parallel

Lecture 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 Parallel：Tensor Parallel（单机之间）以及 Pipeline Parallel（机器之间）

2. 模型装得下后，用 Data Parallem 进行拓展。

3. 如果 Batch Size 太小，此时通信开销会占主导。在本地算多个 batch 的梯度，累加后再统一进行一次 All-Reduce 同步。相当于用计算时间换取通信频率的降低，让通信效率变高。