Deep Learning for Computer Vision——Large Scale Distributed Training

第一章：底层硬件与集群架构

1.1 认识 GPU：NVIDIA H100

• GPU (Graphics Processing Unit)：最初为图形设计，现已成为通用的并行处理器（General Parallel Processor）。

• 物理外观 ：核心计算芯片周围环绕着 HBM 高带宽内存）。H100拥有 80 GB HBM ，核心到内存的带宽高达 3352 GB/sec。

• 流多处理器 (SM)：H100 拥有 132 个 SM。每一个 SM 内部：

  • 内存层次：256 KB L1 缓存 / 寄存器文件。越靠近计算核心，容量越小但速度越快。

  • FP32 核心 ：128个。每个时钟周期执行一个标量运算 a*x + b。整个 SM 每周期可执行 256 FLOPs。

  • Tensor Core（张量核心） ：4个。这是吞吐量的主要来源！ 专门设计用于矩阵乘法。一个张量核心每周期执行 [16x4] @ [4x8] + [16x8]，算力高达 1024 FLOPs。整个 SM 借助 Tensor Core 每周期可执行 4096 FLOPs。

  • 混合精度 ：Tensor Core 通常使用 16-bit 输入（如 BF16/FP16），计算后以 32-bit 输出。⚠️ 初学者注意：在 PyTorch 中如果不使用 .to(torch.bfloat16) 或自动混合精度 (AMP)，模型会跑在普通的 FP32 核心上，速度会慢约 20 倍！

1.2 从单卡到超级计算机集群

• 内存层次结构的延伸：不仅是芯片内部，在集群层面带宽也是逐级递减的（参考图3）。

• Meta Llama3 集群案例：

  • 服务器 (Server)：8× GPU，卡间通信带宽约 900 GB/s。

  • 机架 (Rack)：2台服务器 = 16× GPU。

  • Pod：192个机架 = 3072× GPU，Pod内任意GPU通信约 50 GB/s。

  • 集群 (Cluster) ：8个Pod = 24576× GPU。跨Pod带宽 < 50 GB/s。

  • 💡 核心思想：把整个数据中心当成一台拥有 2.4万个GPU、4.15亿个FP32核心、算力高达 24 ExaFLOPs 的超级计算机！

1.3 其他硬件：Google TPU v5p

• 定制芯片，单芯片 459 TFLOP/s（BF16），95GB 内存，8960个芯片组成一个 Pod。

第二章：四种基本并行策略（4D Parallelism）

面对一个 Transformer 模型，我们有一个形状为 (Batch, Sequence, Dimension, Layers) 的 4D 张量。我们可以在这四个维度上切分并行计算：

1. 数据并行 (DP) ：切分 Batch 维度。

2. 上下文并行 (CP) ：切分 Sequence 维度（处理超长文本）。

3. 流水线并行 (PP) ：切分 Layers 维度（按层切分）。

4. 张量并行 (TP) ：切分 Dimension 维度（切分矩阵乘法）。

第三章：数据并行 (DP) 与完全分片数据并行 (FSDP)

3.1 数据并行 (DP)

• 原理 ：每个 GPU 都保留一个完整的模型副本，但分配不同的数据批次。损失 L=。

• 数学推导 ：因为梯度是线性的，。每个 GPU 计算自己小批次的局部梯度，然后通过 All-Reduce 操作求平均。

• 流程：前向传播 -> 计算各自 Loss -> 反向传播计算局部梯度 -> All-Reduce 平均梯度 -> 所有 GPU 用相同的梯度更新本地模型，保持权重同步。

• ❌ 致命瓶颈 ：模型太大装不下！训练一个 100B 参数的模型，Adam 优化器需要保存权重、梯度、一阶动量、二阶动量（4个数字/参数），每个数字2字节（FP16），单卡需要 800GB 显存！

3.2 完全分片数据并行 (FSDP / ZeRO)

• 原理 ：既然单卡装不下完整模型，就把模型权重也切分！每个权重 只归属于一个 GPU 拥有，其对应的梯度和优化器状态也只由该 GPU 管理。

• FSDP 六步循环：

  1. 前向传播 Layer i ：拥有该层权重的 GPU 将 ​ 广播给所有 GPU。其他 GPU 一边计算，一边预取下一层的权重。

  2. 计算完毕，立即丢弃 ，省显存。

  3. 后向传播 Layer i ：再次广播 。

  4. 所有 GPU 计算局部梯度 ​，计算完立即丢弃 。

  5. 将梯度发送回拥有者 GPU 进行聚合。

  6. 拥有者更新 。

• 优点：显存占用急剧下降，可以训练超出单卡100倍的模型。

• 缺点：通信量增加 50%（前向/反向各广播一次，反向再发送一次）。

3.3 混合分片数据并行 (HSDP)

• 思想：GPU 集群分为 M 组 × K 个 GPU（二维网格）。

• 组内：执行 FSDP（利用 900GB/s 的高带宽）。

• 组间：执行 DP（只同步梯度，利用较慢的 50GB/s 网络）。

• 例子（参考图8）：2组，每组4个GPU。组内高带宽频繁通信，组间低带宽偶尔通信。完美契合硬件拓扑结构。

第四章：显存优化------激活检查点 (Activation Checkpointing)

4.1 激活值的内存瓶颈

以 Llama3-405B 为例：126层，D=16384，序列长度 4096。仅 FFN 中间层的激活值就需要 2 * 126 * (4 * 16384) * 4096 bytes ≈ 63GB。如果加上其他激活值，显存瞬间爆炸。

4.2 激活检查点原理（时间换空间）

• 正常前向+反向：O(N) 计算，O(N) 内存（保存所有层激活用于求导）。

• 全部重计算：O(N²) 计算，O(1) 内存（不保存激活，求导时从头重算，计算量呈平方级，太慢了！）。

• 折中方案（√N 检查点）：

  • 前向过程中，每隔 √N 层保存一个"检查点"（Checkpoint）。

  • 反向传播时，从最近的检查点开始，重新计算该段内未保存的激活值。

  • 复杂度：O(N√N) 计算量，O(√N) 内存占用。这在实践中是完美的平衡点（参考图11右侧的矩阵计算示意图）。

第五章：更复杂的并行策略（CP, PP, TP）

5.1 上下文并行 (Context Parallelism, CP)

• 场景：序列长度超过 128K，单卡显存存不下长序列的注意力矩阵。

• 做法：将序列切分到多个 GPU 上。

• 难点：Self-Attention 是全局的，每个 token 都要看其他所有 token。

• 方案：

  • Ring Attention：将 KV 块在 GPU 之间传递（环形传递），内外层循环计算。复杂但可以支持无限长。

  • Ulysses：不切分注意力矩阵，直接在多头注意力 (Multi-head) 的 Head 维度上切分，每张卡只算几个头，计算完再 All-Gather 拼接。受限于 Head 的数量。

5.2 流水线并行 (Pipeline Parallelism, PP)

• 做法：按层切分，GPU1 算前几层，GPU2 算中间层......将不同层分配到不同 GPU。

• 问题 ：GPU 空闲等待（Bubble）。因为 GPU2 必须等 GPU1 算完才能算，前向时 GPU1 在工作，GPU2 在发呆。最大 MFU 仅为 1/N（N为GPU数）。

• 解决：微批次 (Micro-batches)：把数据切分成多个小批次（如 4 个）。当 GPU1 算完微批次 1 传给 GPU2 时，GPU1 立即开始算微批次 2。这样让流水线满载。

• 效果：4路流水线并行，4个微批次，MFU 从 25% 提升到 57.1%。

5.3 张量并行 (Tensor Parallelism, TP)

• 做法 ：把矩阵乘法分块。例如 X [NxD] @ W [DxD] = Y [NxD]。

• 切分 ：把 W 切分成列 W1, W2, W3, W4（分给4个GPU），每个 GPU 计算 X @ Wi = Yi。最后通过 All-Gather 拼接成完整的 Y。（参考图22）

• 两层情况 ：第一层列切分，第二层行切分（参考图23）。Z = Y1U1 + Y2U2 + ...，最后需要进行 All-Reduce 求和。

• 特点 ：频繁通信，延迟极高。通常限制在单机内（如 8 GPU 节点内）使用，绝不跨节点。

第六章：综合应用与性能评估

6.1 多维并行 (ND Parallelism)

• 做法：同时使用 TP、CP、PP、DP，将 GPU 组织成 4D 网格。

• Llama3-405B 实例配置（参考图24）：

  • 8192 GPU：TP=8, CP=1, PP=16, DP=64，MFU=43%。

  • 16384 GPU：TP=8, CP=16, PP=16, DP=8，MFU=38%。

  • 规律：TP 尽量小（放在机内，带宽高），PP 和 DP 放在机架间，CP 处理长序列。调优的目标是最大化 MFU。

6.2 核心评价指标：HFU 与 MFU

• HFU (Hardware FLOPs Utilization) ：硬件算力利用率。衡量实际执行了多少计算（包括重计算、数据增强、优化器等"无用功"）。H100 理论峰值 989.4 TFLOP/s，实际上大矩阵乘法（Matmul）大概能达到 80% 的 HFU。

• MFU (Model FLOPs Utilization) ：模型算力利用率。这才是我们真正关心的！

  • 计算公式：MFU=理论模型计算量/理论峰值算力实际迭代时间MFU=实际迭代时间理论模型计算量/理论峰值算力​

  • 理论模型计算量（前向+反向） ≈ 6 * N * D （N是Token数，D是模型参数量，反向是前向的2倍）。

  • 行业标准：MFU > 30% 是 Good，> 40% 是 Excellent。

  • 为什么变低了？因为"助手"计算（数据加载、优化器更新、通信）不产生"模型有用计算"，但会消耗时间。

第七章：实战训练指南（Scaling Recipe）

如何选择并行策略？

1. 起步 ：先用 数据并行 (DP)，模型小于 1B 参数，GPU 数量 < 128。尽可能最大化单卡 Batch Size。

2. 模型变大 (>1B) ：切换到 FSDP 或 HSDP。

3. 显存爆了 ：加上 激活检查点 (Activation Checkpointing)。

4. GPU 很多 (>256) ：考虑 HSDP（充分利用组内带宽）。

5. 超大模型 (>50B) 或超长序列 (>16K) ：必须使用 CP, PP, TP 等高级策略。

6. 终极调优 ：不断调整各个并行维度，以最大化 MFU 为最终目标。