大规模分布式计算的重要性

当代深度学习的默认形态是分布式训练，在工业界、创业公司乃至学术界看到的大模型，背后几乎都不是单卡训练出来的，如今的训练经常需要同时动用几十、几百、上千甚至上万张设备，因此需要学习一套与单机/单卡不同的系统方法，既包含硬件，也包含算法与并行策略。

GPU 硬件架构

GPU 里面有什么

GPU 就是 Graphics Processing Unit，最初用于图形渲染，因为其高并行性而被用于通用计算，过去十年内，单卡算力增长了越 1000 倍。
以 NVIDIA H100 为例，内部存在显著的存储层次结构：

Inside a GPU: NVIDIA H100

• HBM（High Bandwidth Memory）：80 GB 容量，带宽为 3 TB/s，距离计算核心较远，因此延迟较高，需要带宽管理。
• L2 Cache：50 MB 容量，位于芯片中部，速度更快。
• SM（Streaming Multiprocessor）：H100 这一代的数量级为 132 SMs，每个 SM 中有更加贴近计算的片上存储与执行单元，例如 L1 Cache、registers、以及用于线程协作与数据复用的 shared memory。

H100 Streaming Multiprocessor

SM 中还包括 FP32 Cores 与 Tensor Cores，FP32 Cores 更像是传统的通用计算流水线，做的是标量或者小向量层面的乘加计算，很灵活但是吞吐量上限相对保守。而 Tensor Cores 则是为矩阵乘法定制的专用单元，一次指令将一次小矩阵块的乘加做完，其算力远超 FP32 Cores。因为训练中大量的层最终都会落到 GEMM（General Matrix Multiply），因此当计算图更加矩阵友好，并且使用混合精度去训练时，能够最大程度使用 Tensor Cores，效率会更高。

GPU 集群：以 llama3 为例

为了训练大模型，必须将成千上万个 GPU 连接起来，系统可以被抽象为一台巨型计算机，其通信带宽随着层级上升而下降。

GPU Cluster

• H100 GPU：3352 GB/sec（inside the GPU）。
• GPU Server：8 卡之间的带宽为900 GB/sec。
• Server Rack：即两个 Server，包含 16 个 GPU。
• Pod：192 个 Rack，共 3072 个 GPU，这些 GPU 之间的带宽为50 GB/sec。
• Cluster：8 个 Pods，共 24576 个 GPU，GPU 之间的带宽来到了< 50GB/sec。

算法的设计需要考虑这种非均匀访问的特性，尽量在带宽高的层级内进行频繁通信，在带宽低的层级间减少通信。

分布式并行策略

为了在硬件上训练巨大模型，需要将计算在四个维度上进行拆分：数据、层、序列以及维度。

数据并行性（Data Parallelism，DP）

将一个更大的 minibatch 拆分为多份，分别交给多张 GPU 各自做一次完整的 forward 和 backward，然后将梯度汇总起来再更新参数，因此 DP 的核心是将同一次迭代的梯度对齐到同一个更新上。

可以将一次迭代的目标写为以下公式：

对参数求梯度：

如果第 $m$ 张 GPU 拿到自己的那份第 $N$ 个样本，它会计算一个本地平均梯度：

那么全局梯度就是：

即如果将 minibatch 扩大到 $MN$，并且分到 $M$ 张 GPU 上，那么梯度对这个平均操作是线性的，于是每张 GPU 只要算自己那一份数据的梯度，最后将梯度做平均，等价于在单卡上对 $MN$ 个样本算一次梯度并且更新。

Data Parallelism

具体的流程如下：

• 先在所有 GPU 上初始化（同一份）权重与优化器状态。
• 每次迭代将一个 minibatch 切成 $M$ 份，每张 GPU 得到 $\frac{B}{M}$。
• 各自做 forward。
• 各自做 backward 得到本地梯度。
• 跨 GPU 进行平均梯度。
• 所有 GPU 使用同一个平均梯度更新各自的本地权重。

其中各自做 backward 得到本地梯度和可以并行（流水线）。DP 在本质上就是计算完全并行，forward 与 backward 在各卡独立完成，几乎不需要在层与层之间传递激活；并且同步发生在梯度处。但是模型的大小仍然被单卡的显存所限制，在纯 DP 中每张 GPU 都要存储一整份模型参数，训练时不仅需要存储权重，还需要存储梯度以及优化器状态。

Fully Sharded Data Parallelism（FSPD）

FSDP 中，每个权重 $W_{i}$ 由某一张 GPU 作为 owner 持有，这个 owner 同时持有该权重对应的梯度与优化器状态。
在一次迭代中，每层执行的步骤大致如下：

1. 前向之前：拥有 $W_{i}$ 的 owner 将 $W_{i}$ 广播给所有 GPU，让大家都拿到这一层的完整权重以便于计算。
2. 所有的 GPU 使用各自的数据分片完成该层的 forward，然将本地的 $W_{i}$ 删除掉。
3. 反向之前 owner 再次将所有的 $W_{i}$ 广播给所有 GPU。
4. 所有 GPU 完成该层 backward，得到本地的 $\frac{\partial L}{\partial W_{i}}$，并且将本地的 $W_{i}$ 删除掉。
5. 反向之后所有 GPU 将各自算到的本地梯度 $\frac{\partial L}{\partial W_{i}}$ 发会给 owner，并且删除掉本地梯度副本。
6. Owner 汇总梯度并且进行该权重的梯度更新，由于 owner 也同时持有优化器状态，因此更新是就地完成的。

FSPD

FSDP 的策略就是通过用完就删来节省显存，前向结束后就删掉 $W_{i}$，反向阶段如果要计算该层的梯度，就必须再次获得 $W_{i}$，以通信换显存。
在显式工程中，有两个较为关键的优化：

• Prefetch：在计算第 $i$ 层 forward 的同时，提前去 fetch $W_{i+1}$，让下一层的权重广播尽量和当前层的计算重叠，从而减少通信裸露在必须等待的路径上。
• 不要立即删除最后一个权重：当做完某曾 forward 时，很快要进入相邻阶段时，如果立即删除权重，马上又要重新再广播一遍，就会造成刚发完又重发的低效抖动，保留一下能够减少无意义的来回传输。

Hybrid Sharded Data Parallel（HSDP）

FSDP 通过权重分片来将单卡显存压下去，但是代价是每一层在 forward/backward 过程中需要反复把权重广播/聚合。FSDP 一次完整 forward + backward 相当于要进行三次网络权重规模的通信（forward 需要传权重，backward 需要传权重并传梯度），而朴素 DP 则只要做一次分配。
GPU 集群的互联是分层的，同一台机器/同一 pod 中互联更快，跨机器/跨机架更慢。HSDP 就是通过将 $N$ 张 GPU 切分成 $M组\times K卡$, 在组内做 FSDP，组间做 DP。

HSDP

这也是一个很套路的系统设计：通信量大的轴（FSDP）尽量绑定到硬件上最快的拓扑层级，通信量小的轴（DP）则跨越更大的拓扑范围。

Activation Checkpointing

即使已经使用了 FSDP/HSDP 将参数、梯度、优化器状态切分到多卡上，训练仍然可能会卡在激活占用的显存上。训练时需要参数，反向传播过程中需要更多显存的常常是中间激活，因为标准反向传播在计算某一层的梯度时，需要用到该层前向产生的激活。前向把 $A_{i}$ 变成 $A_{i+1}$，反向需要 $(A_{i},G_{i+1})$ 来计算 $G_{i}$。

Activation Checkpointing

如果想要让反向传播一次性进行完毕，最朴素的做法是将所有层的激活都保存下来，于是激活相关的内存随着网络深度（层数 $N$）线性增长。因此反向传播过程中，要么使用当时存储下来的激活，要么通过前向传播将当时的激活计算回来。

如果在前向时几乎不存储激活，反向时每需要某层激活，就从更早的 checkpoint 将前向再跑一遍，将其计算回来，这样可以非常节约显存。空间复杂度来到了 $O(1)$，但是计算复杂度则来到了 $O(N^2)$，这显然是不可接受的。

折中的方案就是不要什么都不存，也不要什么都存，而是存储少量的 checkpoints，其余激活在反向阶段按需重算。如果设置 $C$ 个 checkpoints，那么计算复杂度就会变成 $O\left( \frac{N^2}{C} \right)$，而空间复杂度变成 $O(C)$。当有 $C$ 个 checkpoints 时，网络被分割为大约 $C$ 段，每段长度约为 $\frac{N}{C}$。反向经过某一段时，不会为了每一层都从头算到尾，而是从该段段首的 checkpoint 出发，在这一内部将必要的激活重新计算出来，重算的范围就被限制在这一小段中。常见的经验是选取 $\sqrt{ N }$ 级别，此时的计算复杂度为 $O(N\sqrt{ N })$，空间复杂度为 $O(\sqrt{ N })$。

HFU & MFU

此处主要是两套屏幕训练效率的指标体系。HFU 解释单纯将矩阵乘法运算做到多快，而 MFU 则用于解释端到端的训练循环将硬件峰值的占用比例。

Hardware FLOPs Utilization

HFU

HFU 的定义非常直接，以 H100 为例，Tensor Cores 的 16 bit 矩阵乘法的理论峰值是989.4 TFLOP/s，HFU 就是在问实际跑出来的矩阵乘法吞吐占理论峰值的多少比例。因此其局限性也在于其只关注矩阵乘法这一个算子族，不计入 activation checkpointing 带来的额外重计算，也不计入数据增强、优化器、预处理等对迭代时间的影响。

Model FLOPs Utilization

MFU

MFU 是为了修正 HFU 中的只看单算子，不看全流程的缺陷，关注的是 GPU 的理论峰值 FLOPs 中，有多少比例被用于有用的模型计算，并且直接把端到端迭代时间纳入分母，使其成为分布式训练调参时更可信的总目标。其大致的计算流程如下：

1. 估算一次 forward+backward 中 matrix multiply 的总 FLOPs（常用近似：backward ≈ 2× forward；并且忽略非线性、归一化、residual 等 elementwise op，因为它们通常跑在 FP32 cores 上，且不是主要 FLOPs 大头）。
2. 查询设备的理论峰值。
3. 得到一个理论最短时间 $t_{theoretical}=\frac{FLOP_{theoretical}}{sec_{theoretical}}$。
4. 实测真实迭代时间 $t_{actual}$：包含 data loading、forward、backward、optimizer step。
5. 定义 $MFU=\frac{t_{theoretical}}{t_{actual}}$。

Context Parallelism（CP）

Transformer 的计算对象是一段长度为 $S$ 的序列，当想要训练/微调长上下文模型时，瓶颈往往不是参数量本身，而是序列太长导致的激活显存与 attention 的代价，因此 CP 的核心想法是用多张 GPU 共同处理同一长序列，把序列维度切开。

可以按照是否好并行，将Transformer 中的层分割为三类 ：

• LayerNorm/Residual：最容易，因为其没有可学习权重，并且每个 token 的操作基本是局部的，所以沿着序列维度切开后，每张卡自己计算自己的即可，几乎必须要跨卡同步就能够完成前向与反向。
• MLP 与 QKV projection：比较容易，但是也需要像 DP 一样进行梯度同步。MLP 有权重矩阵，但它对不同 token 的计算是独立的，因此把 token 分到不同的 GPU 之后，前向计算天然并行，代价在于每张 GPU 需要保留一份相同的权重副本，反向时再像 DP 那样进行梯度同步。QKV projection 也是一样，其是线性层，对 token 独立，因此也可以按照序列并行以及进行参数梯度同步。
• Attention：最难并行的部分，因为 attention 的每个 query token 需要看见全序列的 key/value，因此把序列分到不同 GPU 之后，如果不做额外的设计，就会立刻遇到信息不全的问题，单卡仅拿到局部 token 的 K/V，无法计算出全局注意力。

Context Parallelism

针对这个问题，可以有两个解决方案：

• Ring Attention：将 attention 相关的计算切成 block 并且分发到各个 GPU，然后在实现上用一种环形的数据流来完成所需的信息交换。即每张 GPU 固定负责一段 query block，然后在一个循环中不断拿到别的 GPU 的 key/value block，每次就计算这一段 Q 对当前这段 K/V 的部分贡献，最终将所有 block 的贡献积累起来得到等价于全序列 attention 的结果，这样的实现较为复杂，但是可以扩展到非常长的序列。

Ring Attention

• Ulysses：不在 GPU 间分布整个 attention 矩阵，而是利用多头注意力的结构，直接沿着 head 维度切分，让不同 GPU 负责不同的 heads，这样实现会简单很多，但是最大的并行度仅仅局限于 head 数量

Pipeline Parallelism

将模型的不同层物理地放置在不同 GPU 上，在 GPU 边界复制层间激活。

Pipeline Parallelism

但是朴素的实现会导致 GPU 出现大量空闲等待时间，例如将一个 L 层网络拆分成 4 段放到 4 张卡上，那么一次前向并不是每张卡都能同时开始计算，第二段必须等待第一段产出的激活，第三段必须等待第二段的激活；反向同理也有依赖链。这会导致 MFU 的上限可以低到 $\frac{1}{N}$。

而解决方案就是将一个 batch 切分成多个小批次同时在流水线中推进，把一个大 batch 切分成多个 microbatch，让不同 microbatch 处在流水线的不同 stage 上，从而将原来的那些等待永别的 microbatch 的计算填满。

Pipeline Parallelism 流水线

Tensor Parallelism

当希望在同一个样本/同一个序列上也能够把计算与显存分摊到多卡时，最直接的做法是沿着 Dim 去切分线性层的权重矩阵，让一次大的矩阵乘法变成多张 GPU 的 block matrix multiply。

Tensor Parallelism

例如有一个经典的线性层 $Y = XW, X \in \mathbb{R}^{N\times D}、 W \in \mathbb{R}^{D \times D}、 Y \in \mathbb{R}^{R\times D}$，在 4-way TP 下将 $W$ 切分为 $[1 \times 4]$ 的列块 $W_{1},W_{2},W_{3},W_{4}$，于是第 $i$ 张 GPU 只需要计算自己的那一块 $Y_{i}=XW_{i}$，输出的 $Y$ 自然也被分割为了 $Y_{1},Y_{2},Y_{3},Y_{4}$。
但是问题就在于，在这一层切分计算完毕之后，是否需要将完整的 $Y$ 拼接起来传给下一层，这样就需要做跨 GPU 的收集，导致在关键路径上，成为计算效率的瓶颈。核心的方法是，当有两层连续的线性层

可以让第一层将 $W$ 按照 $[1 \times 4]$ 切分，第二层 $U$ 按照 $[4 \times 1]$ 的方式切分，这样全局的输出就满足

这样收集就会推迟发生，能够让通信的位置更加可控、也更容易与其他并行维度一起做整体的调度。

ND Parallelism

当模型规模进一步扩大时，分布式训练不再是“DP/TP/PP/CP 选一个”，而是把 GPU 组织成一个多维网格，在不同维度上同时应用不同并行策略，从而同时应对显存容量、计算吞吐与通信拓扑三类约束。一个 Transformer 的训练张量天然可以抽象为 $(\text{Batch}, \text{Sequence}, \text{Dim})$，而模型本身还有层数 $L$，因此最自然的四个切分轴分别对应：DP 切 Batch、CP 切 Sequence、TP 切 Dim、PP 切 Layers；当目标是“训练最大的模型”时，最合理的答案往往就是四种并行都要用，并且要用得彼此匹配。

ND Parallelism 的关键不是概念上的“叠加”，而是把总 GPU 数 $N$ 显式分解成一个乘积

从而让每张 GPU 在训练系统里拥有一个四维坐标：它既属于某个 data-parallel group，也属于某个 pipeline stage，同时还是某个 context slice 与 tensor-parallel group 的成员；这样做的意义在于，并行策略变成了可配置参数，需要决定 $n_{\text{TP}}$ 应该是多少、应该映射到集群拓扑的哪一层、以及它与 $n_{\text{DP}}, n_{\text{PP}}, n_{\text{CP}}$ 的比例关系是否合理。
在这种多维组合里，系统设计的第一原则依然是拓扑感知：把通信量大的维度尽量绑定到更快的互联域，把通信量小的维度才跨越更慢的互联域，从而避免把高频大通信暴露在最差链路上；与此同时，FSDP/HSDP 与 activation checkpointing 这类“用通信换显存、用计算换显存”的手段，会把瓶颈从“放不下”转移为“更难跑满”，因此最终评估标准也必须回到端到端效率。相比只度量矩阵乘法速度的 HFU，ND Parallelism 更强调以 MFU 目标进行闭环调参：通过选择合适的 $(n_{\text{DP}}, n_{\text{PP}}, n_{\text{CP}}, n_{\text{TP}})$ 与拓扑映射，让迭代时间里尽可能多的比例用于有效的模型计算，从而把“很多 GPU”真正转化为“更高吞吐”。

参考如下：

https://cs231n.stanford.edu/index.html