第一章底层硬件与集群架构1.1 认识 GPUNVIDIA H100GPU (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。Pod192个机架 3072× GPUPod内任意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/sBF1695GB 内存8960个芯片组成一个 Pod。第二章四种基本并行策略4D Parallelism面对一个 Transformer 模型我们有一个形状为(Batch, Sequence, Dimension, Layers)的 4D 张量。我们可以在这四个维度上切分并行计算数据并行 (DP)切分Batch维度。上下文并行 (CP)切分Sequence维度处理超长文本。流水线并行 (PP)切分Layers维度按层切分。张量并行 (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 六步循环前向传播 Layer i拥有该层权重的 GPU 将 广播给所有 GPU。其他 GPU 一边计算一边预取下一层的权重。计算完毕立即丢弃省显存。后向传播 Layer i再次广播。所有 GPU 计算局部梯度计算完立即丢弃。将梯度发送回拥有者 GPU 进行聚合。拥有者更新。优点显存占用急剧下降可以训练超出单卡100倍的模型。缺点通信量增加 50%前向/反向各广播一次反向再发送一次。3.3 混合分片数据并行 (HSDP)思想GPU 集群分为 M 组 × K 个 GPU二维网格。组内执行 FSDP利用 900GB/s 的高带宽。组间执行 DP只同步梯度利用较慢的 50GB/s 网络。例子参考图82组每组4个GPU。组内高带宽频繁通信组间低带宽偶尔通信。完美契合硬件拓扑结构。第四章显存优化——激活检查点 (Activation Checkpointing)4.1 激活值的内存瓶颈以 Llama3-405B 为例126层D16384序列长度 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, TP5.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/NN为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 实例配置参考图248192 GPUTP8, CP1, PP16, DP64MFU43%。16384 GPUTP8, CP16, PP16, DP8MFU38%。规律TP 尽量小放在机内带宽高PP 和 DP 放在机架间CP 处理长序列。调优的目标是最大化 MFU。6.2 核心评价指标HFU 与 MFUHFU (Hardware FLOPs Utilization)硬件算力利用率。衡量实际执行了多少计算包括重计算、数据增强、优化器等“无用功”。H100 理论峰值 989.4 TFLOP/s实际上大矩阵乘法Matmul大概能达到80% 的 HFU。MFU (Model FLOPs Utilization)模型算力利用率。这才是我们真正关心的计算公式MFU理论模型计算量/理论峰值算力实际迭代时间MFU实际迭代时间理论模型计算量/理论峰值算力理论模型计算量前向反向 ≈6 * N * DN是Token数D是模型参数量反向是前向的2倍。行业标准MFU 30% 是 Good 40% 是 Excellent。为什么变低了因为“助手”计算数据加载、优化器更新、通信不产生“模型有用计算”但会消耗时间。第七章实战训练指南Scaling Recipe如何选择并行策略起步先用数据并行 (DP)模型小于 1B 参数GPU 数量 128。尽可能最大化单卡 Batch Size。模型变大 (1B)切换到FSDP或HSDP。显存爆了加上激活检查点 (Activation Checkpointing)。GPU 很多 (256)考虑HSDP充分利用组内带宽。超大模型 (50B) 或超长序列 (16K)必须使用CP, PP, TP等高级策略。终极调优不断调整各个并行维度以最大化 MFU为最终目标。
Deep Learning for Computer Vision——Large Scale Distributed Training
第一章底层硬件与集群架构1.1 认识 GPUNVIDIA H100GPU (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。Pod192个机架 3072× GPUPod内任意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/sBF1695GB 内存8960个芯片组成一个 Pod。第二章四种基本并行策略4D Parallelism面对一个 Transformer 模型我们有一个形状为(Batch, Sequence, Dimension, Layers)的 4D 张量。我们可以在这四个维度上切分并行计算数据并行 (DP)切分Batch维度。上下文并行 (CP)切分Sequence维度处理超长文本。流水线并行 (PP)切分Layers维度按层切分。张量并行 (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 六步循环前向传播 Layer i拥有该层权重的 GPU 将 广播给所有 GPU。其他 GPU 一边计算一边预取下一层的权重。计算完毕立即丢弃省显存。后向传播 Layer i再次广播。所有 GPU 计算局部梯度计算完立即丢弃。将梯度发送回拥有者 GPU 进行聚合。拥有者更新。优点显存占用急剧下降可以训练超出单卡100倍的模型。缺点通信量增加 50%前向/反向各广播一次反向再发送一次。3.3 混合分片数据并行 (HSDP)思想GPU 集群分为 M 组 × K 个 GPU二维网格。组内执行 FSDP利用 900GB/s 的高带宽。组间执行 DP只同步梯度利用较慢的 50GB/s 网络。例子参考图82组每组4个GPU。组内高带宽频繁通信组间低带宽偶尔通信。完美契合硬件拓扑结构。第四章显存优化——激活检查点 (Activation Checkpointing)4.1 激活值的内存瓶颈以 Llama3-405B 为例126层D16384序列长度 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, TP5.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/NN为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 实例配置参考图248192 GPUTP8, CP1, PP16, DP64MFU43%。16384 GPUTP8, CP16, PP16, DP8MFU38%。规律TP 尽量小放在机内带宽高PP 和 DP 放在机架间CP 处理长序列。调优的目标是最大化 MFU。6.2 核心评价指标HFU 与 MFUHFU (Hardware FLOPs Utilization)硬件算力利用率。衡量实际执行了多少计算包括重计算、数据增强、优化器等“无用功”。H100 理论峰值 989.4 TFLOP/s实际上大矩阵乘法Matmul大概能达到80% 的 HFU。MFU (Model FLOPs Utilization)模型算力利用率。这才是我们真正关心的计算公式MFU理论模型计算量/理论峰值算力实际迭代时间MFU实际迭代时间理论模型计算量/理论峰值算力理论模型计算量前向反向 ≈6 * N * DN是Token数D是模型参数量反向是前向的2倍。行业标准MFU 30% 是 Good 40% 是 Excellent。为什么变低了因为“助手”计算数据加载、优化器更新、通信不产生“模型有用计算”但会消耗时间。第七章实战训练指南Scaling Recipe如何选择并行策略起步先用数据并行 (DP)模型小于 1B 参数GPU 数量 128。尽可能最大化单卡 Batch Size。模型变大 (1B)切换到FSDP或HSDP。显存爆了加上激活检查点 (Activation Checkpointing)。GPU 很多 (256)考虑HSDP充分利用组内带宽。超大模型 (50B) 或超长序列 (16K)必须使用CP, PP, TP等高级策略。终极调优不断调整各个并行维度以最大化 MFU为最终目标。