计算机视觉-13：大规模分布式训练

GPU 集群

课程以 Llama3-405B 为运行示例，先说明单个 GPU 和大规模集群提供的计算、存储与通信能力。分布式训练并非把单卡程序复制得越多就越快：每张卡完成局部计算后必须交换参数、梯度或激活，计算量通常随 GPU 数近似分摊，通信量却不会按同样比例消失。当一次矩阵乘变短到不足以掩盖数据传输时，新增 GPU 只会增加等待。

GPU 最初服务于图形计算，现在可视为大规模并行处理器。课件以 H100 为例：它配有 80 GB HBM 显存，显存到计算核心的带宽约为 3352 GB/s，并包含约 50 MB 的 L2 缓存和 132 个启用的 Streaming Multiprocessor（SM）。每个 SM 有 256 KB 的 L1 缓存和寄存器空间。

一个 SM 中有 128 个 FP32 核心，每个时钟周期可计算一次 $a\times x+b$ ，即 256 FLOP。它还有 4 个 Tensor Core，可在混合精度下执行矩阵操作 $AX+B$ ；单个 Tensor Core 的示例操作为 $[16\times4][4\times8]+[16\times8]$ ，对应 1024 FLOP，因此一个 SM 可达到 4096 FLOP/周期的 Tensor Core 计算量。

课件中的硬件时间线显示，Tensor Core 峰值吞吐量相较 2013 年早期 GPU 已提高约三个数量级。单卡变快之外，训练还可扩展到多 GPU 集群。

单卡之外，硬件还按服务器、机架和 Pod 组成不同的集群层级。一个服务器可包含 8 张 H100，GPU 之间带宽约为 900 GB/s；两个服务器组成机架，共 16 张 GPU。Llama3 的示例集群中，一个 Pod 由 192 个机架构成，拥有 3072 张 GPU，跨 GPU 通信带宽约为 50 GB/s；8 个 Pod 组成包含 24576 张 GPU 的集群，跨 Pod 通信带宽低于 50 GB/s。

这个集群可被视为一台很大的计算机：总显存约为 1.875 PB，包含约 4.15 亿个 FP32 核心和 1300 万个 Tensor Core，理论峰值约为 24.3 EFLOP/s。不过它不是具有统一共享内存的单台 GPU。节点内 NVLink 远快于跨 Pod 网络，因此一次通信发生在哪一级拓扑上，可能比参加通信的卡数更重要；频繁交换的小张量应留在高速域内，较少发生的大规模同步才适合跨越慢链路。

课程还列出其他训练芯片：TPU v5p 单芯片可提供 459 TFLOP/s 的 BF16 吞吐量和 95 GB 显存，Pod 可达 8960 芯片；AMD MI325X 给出 1300 TFLOP/s BF16 和 256 GB 显存；AWS Trainium2 给出 667 TFLOP/s BF16、96 GB 显存，并以 64 芯片的 UltraServer 组织。

考虑有 $L$ 层的模型，其张量形状写为 $(\mathrm{Batch},\mathrm{Sequence},\mathrm{Dim})$ 。课程按照四个维度切分计算：

策略	切分对象
数据并行（DP）	Batch 维度
上下文并行（CP）	Sequence 维度
流水线并行（PP）	层数 $L$
张量并行（TP）	Dim 维度

四种并行维度

四种并行方式分别沿样本、序列、层和隐藏维切开同一训练计算。选择切分轴的本质是决定“哪部分数据留在本地、哪部分结果必须通信”：数据并行实现最简单但复制整个模型；模型并行能突破单卡容量，却在层内或层间引入更频繁依赖。随后先从通信最少的数据并行开始，再在显存或序列长度无法容纳时逐步引入其他切分。

数据并行

数据并行沿批量维度拆分样本。设每个 GPU 处理 $N$ 个样本，共有 $M$ 个 GPU，则全局小批量包含 $MN$ 个样本：

$\mathcal{L}=\frac{1}{MN}\sum_{m=1}^{M}\sum_{n=1}^{N}\ell(x_{m,n},y_{m,n}).$

损失是所有样本损失的平均，而微分对加法具有线性性，因此可先在每个 GPU 上计算局部平均梯度，再在 $M$ 个 GPU 间平均：

$\nabla\mathcal{L}=\frac{1}{M}\sum_{m=1}^{M} \left(\frac{1}{N}\sum_{n=1}^{N}\nabla\ell(x_{m,n},y_{m,n})\right).$

这说明 GPU 无需交换样本或中间激活，只要在反向末尾汇总梯度，单步结果就与一张无限显存设备处理完整 $MN$ 批量相同。它也是数据并行扩展良好的根源。不过增加 GPU 时若保持每卡批量不变，全局批量会随之增大，单位样本带来的梯度噪声下降，更新次数也可能减少；学习率和训练步数需配套调整，代数等价不保证优化轨迹和泛化完全相同。

普通 DP 的每张 GPU 都持有完整模型和优化器状态，训练流程为：

每张 GPU 加载自己的数据批次。
分别前向传播并计算损失。
分别反向传播，得到本地梯度。
通过 all-reduce 在全部 GPU 间求和并平均梯度。
各 GPU 用相同的平均梯度更新自己的参数。

all-reduce 常拆成 reduce-scatter 与 all-gather：前者一边求和一边把结果分片，后者再让每张卡取得完整平均梯度。反向传播从输出层向输入层逐层产生梯度，因此系统可把已就绪参数装入通信桶，在计算更早层梯度时异步传输；只有通信时间超过可重叠计算时，GPU 才真正停下来等待。随着卡数增加，每卡计算减少而每轮仍需同步近似一份模型大小的梯度，通信最终会成为强扩展瓶颈。

普通 DP 也不解决单卡容量问题，因为每卡仍保存完整模型。按课件的简化估算，每个参数保存参数本身、梯度和 Adam 的两份动量状态，若都按两字节计，10 亿参数约占 8 GB。混合精度训练让大多数矩阵乘和激活使用 BF16/FP16，以更高 Tensor Core 吞吐和更小显存完成前后向，同时可让敏感的累加、优化器状态或主权重保留 FP32，避免小更新被低精度舍入吞掉。FP16 还常需 loss scaling 防止微小梯度下溢，BF16 指数范围更大，通常对此更稳健。混合精度减少的是每个数的成本，没有消除模型状态、激活、通信缓冲区和内存碎片；真实显存必须按各状态的实际类型逐项预算。

显存优化

数据并行中各卡保存完全相同的模型状态，这些副本对计算是必要的，对长期存储却是重复的。全分片数据并行（FSDP）把参数、梯度和优化器状态沿数据并行组切片，每张 GPU 长期只负责约 $1/N$ ；计算某个模块时再临时拼回完整参数。它用额外通信换取近似随卡数下降的模型状态显存，一次典型过程为：

模块前向前，通过 all-gather 汇集各卡的参数分片，临时还原完整参数。
全部 GPU 分别处理自己的数据，完成该模块前向；随后 reshard，释放完整参数副本，并可预取下一模块。
模块反向前再次 all-gather 参数，计算本地参数梯度。
通过 reduce-scatter 对梯度求和，同时只留下每张卡负责的梯度分片。
各 GPU 用本地优化器状态更新自己的参数分片。

all-gather 让每张卡取得当前计算所需的完整参数，reshard 立即释放临时副本；reduce-scatter 同时完成跨卡梯度求和和重新分片，因此各卡只更新自己负责的部分。预取可把下一模块的通信藏在当前模块计算后面，前向末尾也可短暂保留最后模块参数，避免紧接反向时再次收集。FSDP 的思路与 ZeRO-3 接近；ZeRO-1 只去除最大且最少访问的优化器状态副本，ZeRO-2 再分片梯度，ZeRO-3 连参数也分片，显存依次下降，参数 all-gather 的频率与实现复杂度则依次上升。

FSDP 参数分片流程

混合分片数据并行（HSDP）把 $N=M\times K$ 张 GPU 分为 $M$ 个、每组 $K$ 张 GPU。每个组内部运行 FSDP，权重在 $K$ 张 GPU 上分片；各组之间运行 DP。 $K$ 可以达到约百张 GPU。

这种二维组织顺应了网络层级：频繁发生的参数 all-gather 和梯度 reduce-scatter只在高速组内进行；不同组各持有一套分片模型，组间以较低频率同步对应梯度，可跨越相对慢的链路。纯 FSDP 把分片组做得过大虽能继续降低显存，却会让每层集体通信跨越更多慢连接；HSDP 以少量副本换取更局部的高频通信。

参数分片能让单卡放不下的模型继续训练。沿用课件每参数 8 字节的简化口径，1000 亿参数需要 800 GB；分到 80 张 GPU 后，每张只需要约 10 GB 用于参数相关状态。

权重分片后，激活值仍可能占满显存。反向传播计算一层参数梯度时需要该层前向输入，普通训练因而把许多中间结果从前向一直保存到反向；它们随层数、批量和序列长度增长，无法通过参数分片消除。课件以 Llama3-405B 为例：126 层、隐藏维 $D=16384$ 、序列长度 4096，仅 FFN 隐层激活就需要

$2\times126\times(4\times16384)\times4096\ \text{bytes}=63\ \text{GB},$

还未计入其他激活。

将每层写为两个函数：前向计算下一层激活，反向计算上一层梯度，

$A_{i+1}=F_i^{\rightarrow}(A_i),\qquad G_i=F_i^{\leftarrow}(A_i,G_{i+1}).$

课件先用一个全重算模型解释折中：常规前向加反向使用 $O(N)$ 计算和 $O(N)$ 显存；若反向时每次都从输入重算到目标层，显存可降到 $O(1)$ ，计算却会变为 $O(N^2)$ 。在这个教学模型中，每隔若干层保存一次检查点，可得到约 $O(N\sqrt{N})$ 计算和 $O(\sqrt{N})$ 显存的折中。

工程中的激活检查点把网络划分为 $k$ 个连续 segment，只保存各段边界输入；反向进入某段时用边界输入重跑该段前向，重新得到本来应保存的中间值。峰值激活可近似理解为 $O(N/k+k)$ ，这是以额外 FLOP 换显存，特别适合计算资源尚有余量、但更大批量或长序列放不下的场景。段划分、随机状态和算子代价决定实际节省；它只减少为反向保存的激活，无法减少参数、优化器状态或通信缓冲区。

激活检查点与重算

扩展效率

衡量集群扩展要先说明保持什么不变。强扩展固定总模型和全局批量，用更多 GPU 缩短同一步计算；理想时间应降为原来的 $1/G$ ，但每卡计算越来越少，同步延迟最终占据大部分时间。弱扩展随 GPU 数增加总批量或问题规模，使每卡工作量大致不变，更容易维持吞吐，却改变了优化所见的全局批量。若单卡时间为 $T_1$ ， $G$ 卡实测时间为 $T_G$ ，加速比与并行效率可写为 $S_G=T_1/T_G$ 和 $E_G=S_G/G$ ；新增 GPU 带来的价值应看 $E_G$ 是否仍可接受，而不能只看总吞吐是否上升。

课件给出的经验配方是：

约 128 张 GPU、模型约 10 亿参数以内时优先使用数据并行。
每张 GPU 的批量大小尽量填满显存。
模型超过约 10 亿参数时考虑 FSDP。
用激活检查点换取更大的每卡批量。
GPU 数超过约 256 时考虑 HSDP。
GPU 数超过约一千、模型超过 500 亿参数，或序列长度超过 16K 时，再使用 CP、PP 和 TP。

这些阈值只是从“先用通信较少的方案，遇到具体容量瓶颈再切相应维度”推导出的经验顺序，并非模型扩展的定律。模型与数据规模扩大通常能提升能力，训练所需有效 FLOP 也同步增长；并行化只有在新增硬件带来的计算收益大于通信、空泡和重算开销时，才能把这种算法尺度真正转化为可接受的训练时间。配置众多，因此需要统一吞吐指标比较方案。硬件 FLOPs 利用率（HFU）衡量实际矩阵乘吞吐量占设备理论矩阵乘峰值的比例。以 H100 为例，16 位 Tensor Core 矩阵乘理论峰值约为 989.4 TFLOP/s，大矩阵基准中的 HFU 可接近 80%。

HFU 只反映最理想的矩阵乘场景，不计入激活检查点、数据增强、优化器和数据预处理等辅助开销；测量时还需正确同步异步发射的 GPU 内核。

模型 FLOPs 利用率（MFU）衡量设备理论峰值中有多少真正用于模型的有效计算。设一次迭代按模型结构统计得到的有效计算量为 $F_{\mathrm{model}}$ ，GPU 数为 $G$ ，单卡理论峰值为 $P_{\mathrm{peak}}$ ，实测迭代时间为 $t_{\mathrm{actual}}$ ，则

$\mathrm{MFU}=\frac{F_{\mathrm{model}}} {G\,P_{\mathrm{peak}}\,t_{\mathrm{actual}}}.$

分子只统计模型数学上必需的有效计算，分母则是所有 GPU 在实测时间内理论可完成的计算。数据加载、通信等待、流水线空泡和检查点重算都会拉长 $t_{\mathrm{actual}}$ ，却不增加 $F_{\mathrm{model}}$ ，因此 MFU 能把各种隐藏开销统一反映为离峰值还有多远。前向和反向 FLOP 可按网络形状统计，反向常近似为前向两倍。课件将 MFU 高于 30% 视为较好、高于 40% 视为优秀，这是课程示例口径；更晚设备也不必然有更高 MFU，例如计算峰值若比显存带宽增长更快，算子会更早等待数据而非乘法单元。

上下文并行

当单个样本的序列太长，数据并行无法再帮忙，因为这个样本仍需完整放在一张卡上。上下文并行（CP）沿序列切分同一个样本，使每卡只保存一部分 token 激活；它直接解决长上下文容量，但注意力中的每个 query 仍需看到全序列 key/value，因此省下的本地存储会转化为跨卡数据交换。

归一化和残差连接没有参数，沿序列切分后可直接并行。MLP 在每个位置独立运算，但包含权重，因此每张 GPU 保留权重副本，并像 DP 一样同步梯度。QKV 投影与 MLP 类似，也可沿序列并行并同步梯度。

注意力算子最难切分。Ring Attention 让每张 GPU 固定保留本地 query 块，并使 key/value 块沿环依次经过所有卡；每到一块就计算一部分注意力，利用在线 Softmax 维护全局正确的归一化统计。传完一圈后，本地 query 已见过全序列，无需任何卡保存完整 $S\times S$ 矩阵，块通信还能与当前块计算重叠。Ulysses 用 all-to-all 把“各卡有部分序列、全部头”转置成“各卡有全部序列、部分头”，本地完成整头注意力后再转回；逻辑直接，却要求并行度不超过可切分的注意力头数，并对 all-to-all 网络性能敏感。

课件的 Llama3-405B 示例中，第一阶段使用 $S=8192$ ，不采用 CP；第二阶段使用 $S=131072$ ，采用 16 路 CP，使每张 GPU 处理 8192 个 token。

流水线并行

流水线并行（PP）按层切分模型，每张 GPU 只保存连续一段层的参数，激活则像工厂流水线上的工件一样在相邻阶段传递。单个批量必须依次经过各段：第一阶段工作时后面还没收到激活，最后阶段工作时前面又已完成，若一次只处理一个批量， $N$ 路 PP 的最高利用率只有 $1/N$ ，这些首尾空闲区称为 pipeline bubble。

解决方法是把一个批量再拆成多个 microbatch，让它们在各阶段交错执行。对于 $p$ 个阶段、 $m$ 个 microbatch 的 GPipe flush 调度，忽略阶段负载差异时，流水线利用率近似为

$\eta_{\mathrm{pipe}}\approx\frac{m}{m+p-1}.$

式中有效工作对应 $m$ 个 microbatch，填满和排空流水线额外花费 $p-1$ 个阶段时间，所以 $m$ 远大于 $p$ 时固定空泡被摊薄。课件的四路 PP、四个 microbatch 示例将最大 MFU 从 $25\%$ 提升到约 $57.1\%$ 。增大 $m$ 会增加调度和需要管理的激活，阶段计算不均还会由最慢一段决定节拍；1F1B 调度交错前向和反向，可减少同时驻留的激活。

流水线并行

张量并行

张量并行（TP）在单个线性层内部切分权重。对于

$XW=Y,$

可将 $W$ 按列分成 $W_1,\ldots,W_4$ ，每张 GPU 计算 $XW_i=Y_i$ 。每卡只保存和计算一部分输出特征，单个过大的矩阵乘得以跨卡完成；结果 $Y_i$ 仍按特征维分片，若下一算子需要完整 $Y$ 就必须通信收集。

连续两层可交替采用不同的切分方向。第一层按列分片，第二层按行分片：

$Y=[Y_1,Y_2,Y_3,Y_4],\qquad U=\begin{bmatrix}U_1\\U_2\\U_3\\U_4\end{bmatrix}.$

于是每张 GPU 可局部计算一个项，最后再相加：

$Z=Y_1U_1+Y_2U_2+Y_3U_3+Y_4U_4.$

在 Transformer 的 MLP 中，第一层列并行后故意不立刻拼回输出，因为逐元素激活可直接在各分片本地执行；第二层再按行并行，各卡计算上式中的一个加项，最后只做一次 all-reduce。交替切分把两层之间的一次通信消掉。注意力中的 QKV 投影与输出投影也采用相同配对。TP 虽能缩小每卡层内计算和参数，却几乎每层都要同步，通常只能放在 NVLink 或 NVSwitch 等高速域内，否则通信延迟会盖过矩阵乘收益。

张量并行

多维并行

最大规模的训练会同时使用 TP、CP、PP 和 DP，将 GPU 排列为四维网格，并满足

$\mathrm{world\ size}=DP\times CP\times PP\times TP.$

每张 GPU 在四维网格中的坐标决定其通信伙伴。通常让通信最频繁的 TP 留在节点内高速链路，CP 对应长序列组，PP 沿层传递激活，剩余设备再用于通信相对稀疏的 DP。全局批量由 DP 路数、每卡 microbatch 和梯度累积共同决定；盲目提高某一并行度可能只是把容量问题换成通信问题，配置目标是在显存、网络拓扑和优化所需批量之间最大化 MFU。

课件中的 Llama3-405B 示例展示了三种阶段配置：