尧图网页设计 
2. Distributed Training分布式训练2.2 Data Parallelism数据并行2.2.1 数据并行的基本思想2.2.2 每张 GPU 如何完成一次训练2.2.3 Gradient Synchronization 与 All-Reduce2.2.4 Data Parallelism 为什么能加速训练2.2.5 Data Parallelism 的局限2.2.6 Data Parallelism 与 Model Parallelism2.3 Pipeline Parallelism流水线并行2.3.1 Pipeline Parallelism 的基本思想2.3.2 为什么需要 Pipeline Parallelism2.3.3 数据如何在不同 Stage 之间流动2.3.4 Micro-batch让流水线真正并行起来2.3.5 Pipeline Bubble流水线气泡2.3.6 Stage Balance阶段负载均衡2.3.7 Pipeline Parallelism 的核心特点2.4 Tensor Parallelism张量并行2.4.1 Tensor Parallelism 的基本思想2.4.2 为什么需要 Tensor Parallelism2.4.3 Column Parallel 与 Row Parallel2.4.4 Transformer 中如何使用 Tensor Parallelism2.4.5 Tensor Parallelism 的通信开销2.4.6 Tensor Parallelism 的核心特点2.5 三种并行方式的区别2.5.1 Data Parallelism切数据2.5.2 Pipeline Parallelism切 Layer2.5.3 Tensor Parallelism切矩阵2.5.4 DP、PP、TP 对比2.5.5 三种并行方式应该怎么理解2.6 多种并行策略组合2.6.1 DP TP数据并行 张量并行2.6.2 DP PP数据并行 流水线并行2.6.3 TP PP张量并行 流水线并行2.6.4 3D ParallelismDP PP TP2.6.5 GPU 数量如何计算2.6.6 一个 32 GPU 的例子2.6.7 为什么需要组合并行2.6.8 组合并行的代价2. Distributed Training分布式训练2.2 Data Parallelism数据并行每张 GPU 保存一份完整模型将不同 Batch 数据分配到不同 GPU每张 GPU 独立完成 Forward 和 Backward通过 All-Reduce 同步各 GPU 的 Gradient典型实现PyTorch DDP2.2.1 数据并行的基本思想Data Parallelism 的核心是复制模型 切分数据也就是每张 GPU 保存一份完整且相同的模型一个 Global Batch 被切分到不同 GPU每张 GPU 只处理自己分到的数据多张 GPU 同时执行 Forward 和 Backward例如Global Batch Size 1024使用 4 张 GPU每张 GPU 处理 256 条数据因此模型相同数据不同这也是“Data Parallelism”这个名字的来源。2.2.2 每张 GPU 如何完成一次训练假设有两张 GPUGPU 1模型 θ 数据 AGPU 2模型 θ 数据 B两张 GPU 分别执行Forward → Loss → Backward然后得到各自的梯度GPU 1 → g1GPU 2 → g2如果直接使用各自的梯度更新θ1 ← θ - ηg1θ2 ← θ - ηg2两张 GPU 上的模型参数就会逐渐不同。因此 Data Parallelism 还必须进行Gradient Synchronization也就是梯度同步。2.2.3 Gradient Synchronization 与 All-Reduce假设 4 张 GPU 得到g1、g2、g3、g4可以计算平均梯度g (g1 g2 g3 g4) / 4然后所有 GPU 都使用相同的g更新模型θ ← θ - ηg这样下一次训练开始时GPU 1 参数 GPU 2 参数 ... GPU N 参数分布式训练中这种梯度汇总通常通过All-Reduce完成。所以 Data Parallelism 的一次完整 Step 可以记成切分 Batch↓每张 GPU 独立 Forward↓每张 GPU 独立 Backward↓All-Reduce 同步 Gradient↓所有 GPU 同步更新参数2.2.4 Data Parallelism 为什么能加速训练假设一张 GPU 每秒可以处理256 samples使用 4 张 GPU 后每张仍然处理 256 条256 × 4 1024 samples/s因此 Data Parallelism 提升的主要是Training Throughput也就是单位时间能够处理更多训练数据。它并不是让“单条数据”运行得明显更快而是让更多数据同时被处理2.2.5 Data Parallelism 的局限Data Parallelism 的一个重要问题是每张 GPU 都必须保存完整模型。例如模型本身需要40 GB那么使用 4 张 GPU 时仍然是GPU 140 GBGPU 240 GBGPU 340 GBGPU 440 GB而不是40 GB / 4 10 GB因此传统 Data Parallelism可以提高训练吞吐量可以缩短训练时间但不能直接解决“单张 GPU 装不下整个模型”的问题这也是后面需要继续学习Pipeline ParallelismTensor ParallelismZeROFSDP等技术的原因。2.2.6 Data Parallelism 与 Model Parallelism两者最简单的区别Data Parallelism 切数据复制模型Model Parallelism 切模型Data ParallelismGPU 1完整模型 数据 1GPU 2完整模型 数据 2GPU 3完整模型 数据 3Model ParallelismGPU 1模型的一部分GPU 2模型的一部分GPU 3模型的一部分因此 Data Parallelism 最核心的两个关键词就是Replicate Model2.3 Pipeline Parallelism流水线并行将完整模型按照 Layer 切分成多个 Pipeline Stage不同 Stage 分别放到不同 GPU数据按照顺序依次经过多个 Stage使用 Micro-batch 提高多张 GPU 的并行利用率主要问题Pipeline Bubble、Stage 不均衡和 GPU 间通信2.3.1 Pipeline Parallelism 的基本思想Pipeline Parallelism 的核心是切分模型 流水处理数据与 Data Parallelism 不同Data Parallelism 模型复制数据切分而Pipeline Parallelism 模型切分数据依次流过假设一个 Transformer 有 24 层使用 4 张 GPUGPU0Layer 1 ~ 6GPU1Layer 7 ~ 12GPU2Layer 13 ~ 18GPU3Layer 19 ~ 24也就是完整模型 Stage 0 Stage 1 Stage 2 Stage 3每张 GPU 只负责其中一个 Stage。2.3.2 为什么需要 Pipeline ParallelismPipeline Parallelism 最直接解决的问题是模型太大一张 GPU 无法保存完整模型。例如模型大小 120 GB而单张 GPU 显存只有40 GB如果使用普通 Data Parallelism每张 GPU 都要保存完整模型GPU0120 GB ×GPU1120 GB ×GPU2120 GB ×仍然无法训练。而 Pipeline Parallelism 可以把模型拆成多个 StageGPU0Stage 0 ≈ 40 GBGPU1Stage 1 ≈ 40 GBGPU2Stage 2 ≈ 40 GB这样每张 GPU 只保存模型的一部分。因此DP 每张 GPU 保存完整模型PP 每张 GPU 保存部分模型2.3.3 数据如何在不同 Stage 之间流动假设GPU0Layer 1 ~ 6GPU1Layer 7 ~ 12GPU2Layer 13 ~ 18GPU3Layer 19 ~ 24Forward 时Input↓GPU0Stage 0↓Activation A0↓GPU1Stage 1↓Activation A1↓GPU2Stage 2↓Activation A2↓GPU3Stage 3↓Loss因此不同 GPU 之间主要传递的是Activation训练时还需要 BackwardLoss↓GPU3↓Gradient↓GPU2↓GPU1↓GPU0所以可以简单记成ForwardActivation 从前往后传BackwardGradient 从后往前传2.3.4 Micro-batch让流水线真正并行起来如果一次只送一个完整 BatchGPU0 工作↓GPU1 工作↓GPU2 工作↓GPU3 工作那么大部分时间只有一张 GPU 在工作其他 GPU 都在等待。因此 Pipeline Parallelism 通常会把一个 Batch 再切成多个Micro-batch例如Global Batch 64可以切成MB1 16MB2 16MB3 16MB4 16然后让不同 Micro-batch 同时位于不同 StageGPU0MB4GPU1MB3GPU2MB2GPU3MB1这时多张 GPU 就可以同时工作。可以理解成MB1 → Stage0 → Stage1 → Stage2 → Stage3MB2 → Stage0 → Stage1 → Stage2 → Stage3MB3 → Stage0 → Stage1 → Stage2 → Stage3多个 Micro-batch 像工厂产品一样连续进入流水线。2.3.5 Pipeline Bubble流水线气泡即使使用 Micro-batch流水线开始和结束时仍然会出现空闲。例如刚开始Time 1GPU0MB1GPU1IdleGPU2IdleGPU3Idle下一时刻GPU0MB2GPU1MB1GPU2IdleGPU3Idle直到流水线被填满GPU0MB4GPU1MB3GPU2MB2GPU3MB1所有 GPU 才能同时工作。这些等待区域称为Pipeline Bubble通常Micro-batch 数量越多↓Bubble 占比越小↓GPU 利用率越高但 Micro-batch 也不能无限增加还要考虑显存调度开销Batch Size通信开销2.3.6 Stage Balance阶段负载均衡Pipeline Parallelism 不仅要“平均分 Layer”更重要的是让每个 Stage 的计算时间尽量接近例如Stage 010 msStage 110 msStage 240 msStage 310 ms那么 Stage 2 会成为整个 Pipeline 的Bottleneck其他 GPU 即使已经计算完也必须等待 Stage 2。理想情况是T_stage0 ≈ T_stage1 ≈ T_stage2 ≈ T_stage3因此实际划分模型时不一定简单按照每张 GPU 相同 Layer 数还要考虑Attention / MLP 的计算量EmbeddingLM HeadActivation 大小GPU 间通信不同 Stage 的显存占用2.3.7 Pipeline Parallelism 的核心特点Pipeline Parallelism 可以总结为完整模型↓按 Layer 切成多个 Stage↓Stage 分配到不同 GPU↓Batch 再切成多个 Micro-batch↓不同 Micro-batch 同时位于不同 Stage↓形成流水线它的优势是降低单张 GPU 保存模型参数的压力可以训练单卡无法容纳的大模型多个 Stage 可以同时计算不同 Micro-batch主要问题是Pipeline BubbleStage BalanceActivation / Gradient 通信调度更加复杂最后可以记成DP 切数据PP 切 Layer而下一节TP 连一个 Layer 内部的大矩阵都切开2.4 Tensor Parallelism张量并行将同一个 Layer 内部的大矩阵切分到多张 GPU多张 GPU 共同完成同一层的计算常用于 Attention 和 MLP 的大矩阵乘法可以降低单张 GPU 的参数与计算压力代价是需要更频繁的 GPU 间通信典型实现Megatron-LM2.4.1 Tensor Parallelism 的基本思想Tensor Parallelism 的核心是同一层内部切分矩阵假设一个线性层Y XW其中X[b,s,h]W[h,4h]如果权重矩阵W很大可以把它拆成多个部分W [W1 | W2 | ... | WN]然后分别放到不同 GPUGPU0 → W1GPU1 → W2GPU2 → W3...所有 GPU 使用同一个输入X同时计算自己负责的矩阵部分。所以Pipeline Parallelism 不同 GPU 负责不同 Layer而Tensor Parallelism 多个 GPU 共同完成同一个 Layer可以简单记成PP 层与层之间切TP 一层内部切2.4.2 为什么需要 Tensor ParallelismPipeline Parallelism 已经可以把不同 Layer 放到不同 GPU。但是还可能出现一个问题单个 Transformer Layer 本身就很大一张 GPU 仍然放不下。例如 MLP 第一层W1[h,4h]当h 8192时W1[8192,32768]这个矩阵本身就包含大量参数。Attention 中同样存在WqWkWvWo这些大型矩阵。因此可以进一步把一个 Layer 内部的矩阵拆开↓多张 GPU 共同保存↓多张 GPU 同时计算这就是 Tensor Parallelism。所以它主要解决单层模型太大和单层矩阵计算量太大的问题。2.4.3 Column Parallel 与 Row ParallelTensor Parallelism 最常见的矩阵切分方式可以先理解成Column Parallel和Row ParallelColumn Parallel按输出维度切分假设Y XW其中W[h,4h]使用两张 GPU 时可以把 W 按列切开W [W1 | W2]其中W1[h,2h]W2[h,2h]两张 GPU 同时计算GPU0Y1 XW1GPU1Y2 XW2最后Y Concat(Y1,Y2)因此 Column Parallel 可以理解成每张 GPU 负责一部分输出维度Row Parallel按输入维度切分另外一种方式是把权重矩阵按照输入维度切开。例如W [W1 ; W2]输入也对应拆开X [X1 | X2]两张 GPU 分别计算GPU0Y1 X1W1GPU1Y2 X2W2最后需要Y Y1 Y2所以通常需要进行All-Reduce或者类似的归约通信。因此可以简单记成Column Parallel → 输出切开最后 ConcatRow Parallel → 输入切开最后 Sum / Reduce这两种方式组合起来就可以高效拆分 Transformer 中的大型线性层。2.4.4 Transformer 中如何使用 Tensor ParallelismTransformer 一层主要包含Self-Attention和MLP这两部分都有大量矩阵乘法因此非常适合 Tensor Parallelism。Attention 部分Attention 中有Wq、Wk、Wv、Wo还包含多个 Attention Head。例如32 个 Attention Heads使用 4 张 GPUGPU0Head 1 ~ 8GPU1Head 9 ~ 16GPU2Head 17 ~ 24GPU3Head 25 ~ 32每张 GPU 只负责自己的一部分 Attention Head。因此同一层 Attention↓多张 GPU 同时计算↓再通过通信合并结果MLP 部分经典 MLPh → 4h → h第一层W1[h,4h]可以使用 Column ParallelW1 → 按列切分第二层W2[4h,h]可以使用 Row ParallelW2 → 按行切分这样两层可以自然组合Column Parallel↓各 GPU 得到部分中间特征↓Row Parallel↓All-Reduce↓恢复完整输出这也是 Megatron-LM 中非常经典的 Tensor Parallel 思路。2.4.5 Tensor Parallelism 的通信开销Tensor Parallelism 可以把参数和计算拆到多张 GPU但代价是GPU 之间需要频繁通信常见通信操作包括All-ReduceAll-GatherReduce-Scatter例如 Row Parallel 中GPU0 → Y0GPU1 → Y1GPU2 → Y2GPU3 → Y3最终需要得到Y Y0 Y1 Y2 Y3这就需要进行All-Reduce所以 Tensor Parallelism 对 GPU 间互联速度非常敏感。通常更加依赖NVLinkNVSwitch等高速互联。如果计算时间 通信时间那么 GPU 大量时间都会花在等待通信上并行收益就会明显下降。2.4.6 Tensor Parallelism 的核心特点Tensor Parallelism 可以总结成一个 Transformer Layer↓大型权重矩阵切分↓不同 GPU 保存不同矩阵分片↓所有 GPU 同时完成这一层的部分计算↓通过 All-Reduce / All-Gather 等方式交换结果↓得到完整 Layer 输出它的主要优势是降低单张 GPU 保存单层参数的压力分摊大型矩阵乘法的计算量可以让多张 GPU 共同完成同一个 Transformer Layer特别适合 Attention 和 MLP主要问题是GPU 间通信频繁对 NVLink / NVSwitch 等高速互联依赖较强Tensor Parallel Degree 越大通信成本通常越明显最后可以把三种基本并行方式先记成DP 切数据PP 切 LayerTP 切 Layer 内部的矩阵2.5 三种并行方式的区别前面已经分别介绍了Data ParallelismPipeline ParallelismTensor Parallelism三种方法最大的区别就是DP 切数据PP 切 LayerTP 切 Layer 内部的矩阵它们解决的问题并不完全相同。2.5.1 Data Parallelism切数据Data Parallelism 的核心是模型复制 数据切分假设有 4 张 GPUGPU0完整模型 Data0GPU1完整模型 Data1GPU2完整模型 Data2GPU3完整模型 Data3每张 GPU 独立执行Forward → Backward然后通过All-Reduce同步梯度。因此 DP 主要解决单张 GPU 计算速度有限希望多张 GPU 同时处理更多数据。它主要提升的是Training Throughput但是缺点是每张 GPU 都要保存完整模型所以如果模型本身单卡放不下单纯使用 DP 仍然无法训练。2.5.2 Pipeline Parallelism切 LayerPipeline Parallelism 的核心是把模型按照 Layer / Stage 切开例如GPU0Layer 1 ~ 6GPU1Layer 7 ~ 12GPU2Layer 13 ~ 18GPU3Layer 19 ~ 24同一份数据依次流过GPU0 → GPU1 → GPU2 → GPU3因此 PP 主要解决完整模型太大一张 GPU 无法容纳。为了提高 GPU 利用率还需要Micro-batch让不同 Micro-batch 同时位于不同 Pipeline Stage。主要问题是Pipeline BubbleStage BalanceActivation / Gradient 通信2.5.3 Tensor Parallelism切矩阵Tensor Parallelism 切得更细。它不是把不同 Layer 分给不同 GPU而是多个 GPU 共同完成同一个 Layer例如线性层Y XW可以把W切成W1、W2、W3、W4然后GPU0 → W1GPU1 → W2GPU2 → W3GPU3 → W4多张 GPU 同时完成同一个矩阵乘法的一部分。因此 TP 主要解决单个 Transformer Layer 或大型矩阵本身就太大。它经常用于Attention HeadsQ / K / V / WoMLP 大矩阵主要问题是GPU 间通信非常频繁因此通常需要NVLink / NVSwitch等高速互联。2.5.4 DP、PP、TP 对比并行方式切分对象每张 GPU 保存什么数据如何流动主要解决的问题主要代价Data Parallelism数据完整模型不同 GPU 处理不同 Batch提高训练吞吐量梯度同步Pipeline ParallelismLayer / Stage部分模型层数据依次经过多个 Stage模型单卡放不下Pipeline BubbleTensor Parallelism单层内部矩阵矩阵的一部分多 GPU 同时算同一层单层本身太大高频 GPU 通信可以简单记成DP模型不切数据切PP模型按层切TP一层内部继续切2.5.5 三种并行方式应该怎么理解如果模型能够完整放进一张 GPU但训练太慢优先考虑 DP如果整个模型太大一张 GPU 放不下考虑 PP如果连单个 Transformer Layer 都很大考虑 TP但是超大模型训练中通常不会只选择一种。例如DP PPDP TPPP TP甚至DP PP TP一起使用。因此三种并行方式更准确的关系不是谁替代谁而是它们解决不同维度的问题可以组合使用这就自然进入下一节2.6 多种并行策略组合也就是大模型训练中常说的3D Parallelism2.6 多种并行策略组合实际的大模型训练中通常不会只使用一种并行方式。因为DP、PP、TP分别解决的是不同问题DP切数据提高训练吞吐量PP切模型层解决整个模型单卡放不下TP切单层矩阵解决单个 Layer 本身太大因此实际训练中经常组合使用DP TPDP PPTP PP以及DP PP TP其中DP PP TP通常称为3D Parallelism2.6.1 DP TP数据并行 张量并行假设有 8 张 GPU。可以划分成两个 Data Parallel GroupDP Group 0GPU0 ~ GPU3DP Group 1GPU4 ~ GPU7每个 DP Group 内部再使用 Tensor ParallelismGPU0 ~ GPU3共同完成一个模型副本中的矩阵计算。而GPU4 ~ GPU7共同完成另一份相同的模型副本。两组处理不同的数据DP Group 0 → Data ADP Group 1 → Data B因此TP负责把同一个 Layer 内部的大矩阵切开而DP负责复制多份 TP 模型处理不同数据这种组合适合模型单层较大同时还希望通过多份模型副本提高训练吞吐量。2.6.2 DP PP数据并行 流水线并行假设一个模型被划分成 4 个 Pipeline StageStage 0Stage 1Stage 2Stage 3使用 4 张 GPUGPU0 → Stage 0GPU1 → Stage 1GPU2 → Stage 2GPU3 → Stage 3这 4 张 GPU 共同组成一个完整 Pipeline如果还有另外 4 张 GPUGPU4 ~ GPU7就可以再复制一套相同的 Pipeline。于是Pipeline 0 → Data APipeline 1 → Data B两套 Pipeline 最后再进行梯度同步。因此PP负责把一个完整模型按 Layer 切开而DP负责复制多套完整 Pipeline这种组合适合整个模型太大同时还希望通过多份 Pipeline 提高整体训练吞吐量。2.6.3 TP PP张量并行 流水线并行如果模型继续增大可能出现两个问题整个模型太大同时单个 Layer 也很大这时可以同时使用TP PP先使用 Pipeline Parallelism模型 → Stage 0 Stage 1 Stage 2 Stage 3再在每个 Stage 内部使用 Tensor Parallelism。例如 8 张 GPUStage 0 → GPU0 GPU1Stage 1 → GPU2 GPU3Stage 2 → GPU4 GPU5Stage 3 → GPU6 GPU7这里PP 切不同 Layer而TP 切每个 Layer 内部的大矩阵因此可以进一步降低单张 GPU 的模型参数和计算压力2.6.4 3D ParallelismDP PP TP超大模型训练通常会同时使用Data ParallelismPipeline ParallelismTensor Parallelism这就是3D Parallelism三个维度分别负责DP → 数据维度PP → Layer 维度TP → 单层矩阵维度可以把它理解成一个训练集群↓先划分多个Data Parallel Group↓每个 DP Group 内部划分Pipeline Stage↓每个 Pipeline Stage 内部继续划分Tensor Parallel Group所以三种并行方式不是互相替代而是在不同维度共同拆分训练任务。2.6.5 GPU 数量如何计算假设DP 4PP 2TP 4那么总 GPU 数量GPU Total DP × PP × TP也就是4 × 2 × 4 32因此GPU 总数 DP × PP × TP这个公式非常重要。其中TP 4一个 Layer 内部由 4 张 GPU 共同计算PP 2一个模型被切成 2 个 Pipeline StageDP 4完整模型结构一共复制 4 份2.6.6 一个 32 GPU 的例子假设GPU Total 32设置TP 4PP 2那么一份完整模型需要TP × PP也就是4 × 2 8 张 GPU32 张 GPU 一共可以组成32 / 8 4份完整模型。所以DP 4最终DP 4PP 2TP 4并满足32 4 × 2 × 4可以理解成4 个 Data Parallel Group每个 Group 内2 个 Pipeline Stage每个 Stage 内4 张 GPU 做 Tensor Parallel2.6.7 为什么需要组合并行三种并行方式分别解决DP↓训练吞吐量不足PP↓整个模型单卡放不下TP↓单个 Layer 本身太大因此实际大模型训练通常不是DP / PP / TP 三选一而是根据模型大小和硬件条件进行组合最终目标是同时解决模型放不下训练速度太慢GPU 利用率不高的问题。2.6.8 组合并行的代价并行方式越多系统也越复杂。需要同时考虑DP 的 Gradient SynchronizationPP 的 Pipeline BubbleTP 的高频 GPU CommunicationParallel Group 的划分GPU 间通信拓扑NVLink / NVSwitch / InfiniBand显存分配计算负载均衡所以GPU 数量更多 ≠ 一定训练更快真正需要平衡的是计算显存通信并行效率因此这一部分可以最终压缩成DP 切数据PP 切 LayerTP 切矩阵进一步DP TPDP PPTP PPDP PP TP 3D Parallelism并且GPU Total DP × PP × TP这也是后面继续学习DeepSpeed和Megatron-LM的重要基础。Split Batch 大模型训练与推理优化（二）
2. Distributed Training分布式训练2.2 Data Parallelism数据并行2.2.1 数据并行的基本思想2.2.2 每张 GPU 如何完成一次训练2.2.3 Gradient Synchronization 与 All-Reduce2.2.4 Data Parallelism 为什么能加速训练2.2.5 Data Parallelism 的局限2.2.6 Data Parallelism 与 Model Parallelism2.3 Pipeline Parallelism流水线并行2.3.1 Pipeline Parallelism 的基本思想2.3.2 为什么需要 Pipeline Parallelism2.3.3 数据如何在不同 Stage 之间流动2.3.4 Micro-batch让流水线真正并行起来2.3.5 Pipeline Bubble流水线气泡2.3.6 Stage Balance阶段负载均衡2.3.7 Pipeline Parallelism 的核心特点2.4 Tensor Parallelism张量并行2.4.1 Tensor Parallelism 的基本思想2.4.2 为什么需要 Tensor Parallelism2.4.3 Column Parallel 与 Row Parallel2.4.4 Transformer 中如何使用 Tensor Parallelism2.4.5 Tensor Parallelism 的通信开销2.4.6 Tensor Parallelism 的核心特点2.5 三种并行方式的区别2.5.1 Data Parallelism切数据2.5.2 Pipeline Parallelism切 Layer2.5.3 Tensor Parallelism切矩阵2.5.4 DP、PP、TP 对比2.5.5 三种并行方式应该怎么理解2.6 多种并行策略组合2.6.1 DP TP数据并行 张量并行2.6.2 DP PP数据并行 流水线并行2.6.3 TP PP张量并行 流水线并行2.6.4 3D ParallelismDP PP TP2.6.5 GPU 数量如何计算2.6.6 一个 32 GPU 的例子2.6.7 为什么需要组合并行2.6.8 组合并行的代价2. Distributed Training分布式训练2.2 Data Parallelism数据并行每张 GPU 保存一份完整模型将不同 Batch 数据分配到不同 GPU每张 GPU 独立完成 Forward 和 Backward通过 All-Reduce 同步各 GPU 的 Gradient典型实现PyTorch DDP2.2.1 数据并行的基本思想Data Parallelism 的核心是复制模型 切分数据也就是每张 GPU 保存一份完整且相同的模型一个 Global Batch 被切分到不同 GPU每张 GPU 只处理自己分到的数据多张 GPU 同时执行 Forward 和 Backward例如Global Batch Size 1024使用 4 张 GPU每张 GPU 处理 256 条数据因此模型相同数据不同这也是“Data Parallelism”这个名字的来源。2.2.2 每张 GPU 如何完成一次训练假设有两张 GPUGPU 1模型 θ 数据 AGPU 2模型 θ 数据 B两张 GPU 分别执行Forward → Loss → Backward然后得到各自的梯度GPU 1 → g1GPU 2 → g2如果直接使用各自的梯度更新θ1 ← θ - ηg1θ2 ← θ - ηg2两张 GPU 上的模型参数就会逐渐不同。因此 Data Parallelism 还必须进行Gradient Synchronization也就是梯度同步。2.2.3 Gradient Synchronization 与 All-Reduce假设 4 张 GPU 得到g1、g2、g3、g4可以计算平均梯度g (g1 g2 g3 g4) / 4然后所有 GPU 都使用相同的g更新模型θ ← θ - ηg这样下一次训练开始时GPU 1 参数 GPU 2 参数 ... GPU N 参数分布式训练中这种梯度汇总通常通过All-Reduce完成。所以 Data Parallelism 的一次完整 Step 可以记成切分 Batch↓每张 GPU 独立 Forward↓每张 GPU 独立 Backward↓All-Reduce 同步 Gradient↓所有 GPU 同步更新参数2.2.4 Data Parallelism 为什么能加速训练假设一张 GPU 每秒可以处理256 samples使用 4 张 GPU 后每张仍然处理 256 条256 × 4 1024 samples/s因此 Data Parallelism 提升的主要是Training Throughput也就是单位时间能够处理更多训练数据。它并不是让“单条数据”运行得明显更快而是让更多数据同时被处理2.2.5 Data Parallelism 的局限Data Parallelism 的一个重要问题是每张 GPU 都必须保存完整模型。例如模型本身需要40 GB那么使用 4 张 GPU 时仍然是GPU 140 GBGPU 240 GBGPU 340 GBGPU 440 GB而不是40 GB / 4 10 GB因此传统 Data Parallelism可以提高训练吞吐量可以缩短训练时间但不能直接解决“单张 GPU 装不下整个模型”的问题这也是后面需要继续学习Pipeline ParallelismTensor ParallelismZeROFSDP等技术的原因。2.2.6 Data Parallelism 与 Model Parallelism两者最简单的区别Data Parallelism 切数据复制模型Model Parallelism 切模型Data ParallelismGPU 1完整模型 数据 1GPU 2完整模型 数据 2GPU 3完整模型 数据 3Model ParallelismGPU 1模型的一部分GPU 2模型的一部分GPU 3模型的一部分因此 Data Parallelism 最核心的两个关键词就是Replicate Model2.3 Pipeline Parallelism流水线并行将完整模型按照 Layer 切分成多个 Pipeline Stage不同 Stage 分别放到不同 GPU数据按照顺序依次经过多个 Stage使用 Micro-batch 提高多张 GPU 的并行利用率主要问题Pipeline Bubble、Stage 不均衡和 GPU 间通信2.3.1 Pipeline Parallelism 的基本思想Pipeline Parallelism 的核心是切分模型 流水处理数据与 Data Parallelism 不同Data Parallelism 模型复制数据切分而Pipeline Parallelism 模型切分数据依次流过假设一个 Transformer 有 24 层使用 4 张 GPUGPU0Layer 1 ~ 6GPU1Layer 7 ~ 12GPU2Layer 13 ~ 18GPU3Layer 19 ~ 24也就是完整模型 Stage 0 Stage 1 Stage 2 Stage 3每张 GPU 只负责其中一个 Stage。2.3.2 为什么需要 Pipeline ParallelismPipeline Parallelism 最直接解决的问题是模型太大一张 GPU 无法保存完整模型。例如模型大小 120 GB而单张 GPU 显存只有40 GB如果使用普通 Data Parallelism每张 GPU 都要保存完整模型GPU0120 GB ×GPU1120 GB ×GPU2120 GB ×仍然无法训练。而 Pipeline Parallelism 可以把模型拆成多个 StageGPU0Stage 0 ≈ 40 GBGPU1Stage 1 ≈ 40 GBGPU2Stage 2 ≈ 40 GB这样每张 GPU 只保存模型的一部分。因此DP 每张 GPU 保存完整模型PP 每张 GPU 保存部分模型2.3.3 数据如何在不同 Stage 之间流动假设GPU0Layer 1 ~ 6GPU1Layer 7 ~ 12GPU2Layer 13 ~ 18GPU3Layer 19 ~ 24Forward 时Input↓GPU0Stage 0↓Activation A0↓GPU1Stage 1↓Activation A1↓GPU2Stage 2↓Activation A2↓GPU3Stage 3↓Loss因此不同 GPU 之间主要传递的是Activation训练时还需要 BackwardLoss↓GPU3↓Gradient↓GPU2↓GPU1↓GPU0所以可以简单记成ForwardActivation 从前往后传BackwardGradient 从后往前传2.3.4 Micro-batch让流水线真正并行起来如果一次只送一个完整 BatchGPU0 工作↓GPU1 工作↓GPU2 工作↓GPU3 工作那么大部分时间只有一张 GPU 在工作其他 GPU 都在等待。因此 Pipeline Parallelism 通常会把一个 Batch 再切成多个Micro-batch例如Global Batch 64可以切成MB1 16MB2 16MB3 16MB4 16然后让不同 Micro-batch 同时位于不同 StageGPU0MB4GPU1MB3GPU2MB2GPU3MB1这时多张 GPU 就可以同时工作。可以理解成MB1 → Stage0 → Stage1 → Stage2 → Stage3MB2 → Stage0 → Stage1 → Stage2 → Stage3MB3 → Stage0 → Stage1 → Stage2 → Stage3多个 Micro-batch 像工厂产品一样连续进入流水线。2.3.5 Pipeline Bubble流水线气泡即使使用 Micro-batch流水线开始和结束时仍然会出现空闲。例如刚开始Time 1GPU0MB1GPU1IdleGPU2IdleGPU3Idle下一时刻GPU0MB2GPU1MB1GPU2IdleGPU3Idle直到流水线被填满GPU0MB4GPU1MB3GPU2MB2GPU3MB1所有 GPU 才能同时工作。这些等待区域称为Pipeline Bubble通常Micro-batch 数量越多↓Bubble 占比越小↓GPU 利用率越高但 Micro-batch 也不能无限增加还要考虑显存调度开销Batch Size通信开销2.3.6 Stage Balance阶段负载均衡Pipeline Parallelism 不仅要“平均分 Layer”更重要的是让每个 Stage 的计算时间尽量接近例如Stage 010 msStage 110 msStage 240 msStage 310 ms那么 Stage 2 会成为整个 Pipeline 的Bottleneck其他 GPU 即使已经计算完也必须等待 Stage 2。理想情况是T_stage0 ≈ T_stage1 ≈ T_stage2 ≈ T_stage3因此实际划分模型时不一定简单按照每张 GPU 相同 Layer 数还要考虑Attention / MLP 的计算量EmbeddingLM HeadActivation 大小GPU 间通信不同 Stage 的显存占用2.3.7 Pipeline Parallelism 的核心特点Pipeline Parallelism 可以总结为完整模型↓按 Layer 切成多个 Stage↓Stage 分配到不同 GPU↓Batch 再切成多个 Micro-batch↓不同 Micro-batch 同时位于不同 Stage↓形成流水线它的优势是降低单张 GPU 保存模型参数的压力可以训练单卡无法容纳的大模型多个 Stage 可以同时计算不同 Micro-batch主要问题是Pipeline BubbleStage BalanceActivation / Gradient 通信调度更加复杂最后可以记成DP 切数据PP 切 Layer而下一节TP 连一个 Layer 内部的大矩阵都切开2.4 Tensor Parallelism张量并行将同一个 Layer 内部的大矩阵切分到多张 GPU多张 GPU 共同完成同一层的计算常用于 Attention 和 MLP 的大矩阵乘法可以降低单张 GPU 的参数与计算压力代价是需要更频繁的 GPU 间通信典型实现Megatron-LM2.4.1 Tensor Parallelism 的基本思想Tensor Parallelism 的核心是同一层内部切分矩阵假设一个线性层Y XW其中X[b,s,h]W[h,4h]如果权重矩阵W很大可以把它拆成多个部分W [W1 | W2 | ... | WN]然后分别放到不同 GPUGPU0 → W1GPU1 → W2GPU2 → W3...所有 GPU 使用同一个输入X同时计算自己负责的矩阵部分。所以Pipeline Parallelism 不同 GPU 负责不同 Layer而Tensor Parallelism 多个 GPU 共同完成同一个 Layer可以简单记成PP 层与层之间切TP 一层内部切2.4.2 为什么需要 Tensor ParallelismPipeline Parallelism 已经可以把不同 Layer 放到不同 GPU。但是还可能出现一个问题单个 Transformer Layer 本身就很大一张 GPU 仍然放不下。例如 MLP 第一层W1[h,4h]当h 8192时W1[8192,32768]这个矩阵本身就包含大量参数。Attention 中同样存在WqWkWvWo这些大型矩阵。因此可以进一步把一个 Layer 内部的矩阵拆开↓多张 GPU 共同保存↓多张 GPU 同时计算这就是 Tensor Parallelism。所以它主要解决单层模型太大和单层矩阵计算量太大的问题。2.4.3 Column Parallel 与 Row ParallelTensor Parallelism 最常见的矩阵切分方式可以先理解成Column Parallel和Row ParallelColumn Parallel按输出维度切分假设Y XW其中W[h,4h]使用两张 GPU 时可以把 W 按列切开W [W1 | W2]其中W1[h,2h]W2[h,2h]两张 GPU 同时计算GPU0Y1 XW1GPU1Y2 XW2最后Y Concat(Y1,Y2)因此 Column Parallel 可以理解成每张 GPU 负责一部分输出维度Row Parallel按输入维度切分另外一种方式是把权重矩阵按照输入维度切开。例如W [W1 ; W2]输入也对应拆开X [X1 | X2]两张 GPU 分别计算GPU0Y1 X1W1GPU1Y2 X2W2最后需要Y Y1 Y2所以通常需要进行All-Reduce或者类似的归约通信。因此可以简单记成Column Parallel → 输出切开最后 ConcatRow Parallel → 输入切开最后 Sum / Reduce这两种方式组合起来就可以高效拆分 Transformer 中的大型线性层。2.4.4 Transformer 中如何使用 Tensor ParallelismTransformer 一层主要包含Self-Attention和MLP这两部分都有大量矩阵乘法因此非常适合 Tensor Parallelism。Attention 部分Attention 中有Wq、Wk、Wv、Wo还包含多个 Attention Head。例如32 个 Attention Heads使用 4 张 GPUGPU0Head 1 ~ 8GPU1Head 9 ~ 16GPU2Head 17 ~ 24GPU3Head 25 ~ 32每张 GPU 只负责自己的一部分 Attention Head。因此同一层 Attention↓多张 GPU 同时计算↓再通过通信合并结果MLP 部分经典 MLPh → 4h → h第一层W1[h,4h]可以使用 Column ParallelW1 → 按列切分第二层W2[4h,h]可以使用 Row ParallelW2 → 按行切分这样两层可以自然组合Column Parallel↓各 GPU 得到部分中间特征↓Row Parallel↓All-Reduce↓恢复完整输出这也是 Megatron-LM 中非常经典的 Tensor Parallel 思路。2.4.5 Tensor Parallelism 的通信开销Tensor Parallelism 可以把参数和计算拆到多张 GPU但代价是GPU 之间需要频繁通信常见通信操作包括All-ReduceAll-GatherReduce-Scatter例如 Row Parallel 中GPU0 → Y0GPU1 → Y1GPU2 → Y2GPU3 → Y3最终需要得到Y Y0 Y1 Y2 Y3这就需要进行All-Reduce所以 Tensor Parallelism 对 GPU 间互联速度非常敏感。通常更加依赖NVLinkNVSwitch等高速互联。如果计算时间 通信时间那么 GPU 大量时间都会花在等待通信上并行收益就会明显下降。2.4.6 Tensor Parallelism 的核心特点Tensor Parallelism 可以总结成一个 Transformer Layer↓大型权重矩阵切分↓不同 GPU 保存不同矩阵分片↓所有 GPU 同时完成这一层的部分计算↓通过 All-Reduce / All-Gather 等方式交换结果↓得到完整 Layer 输出它的主要优势是降低单张 GPU 保存单层参数的压力分摊大型矩阵乘法的计算量可以让多张 GPU 共同完成同一个 Transformer Layer特别适合 Attention 和 MLP主要问题是GPU 间通信频繁对 NVLink / NVSwitch 等高速互联依赖较强Tensor Parallel Degree 越大通信成本通常越明显最后可以把三种基本并行方式先记成DP 切数据PP 切 LayerTP 切 Layer 内部的矩阵2.5 三种并行方式的区别前面已经分别介绍了Data ParallelismPipeline ParallelismTensor Parallelism三种方法最大的区别就是DP 切数据PP 切 LayerTP 切 Layer 内部的矩阵它们解决的问题并不完全相同。2.5.1 Data Parallelism切数据Data Parallelism 的核心是模型复制 数据切分假设有 4 张 GPUGPU0完整模型 Data0GPU1完整模型 Data1GPU2完整模型 Data2GPU3完整模型 Data3每张 GPU 独立执行Forward → Backward然后通过All-Reduce同步梯度。因此 DP 主要解决单张 GPU 计算速度有限希望多张 GPU 同时处理更多数据。它主要提升的是Training Throughput但是缺点是每张 GPU 都要保存完整模型所以如果模型本身单卡放不下单纯使用 DP 仍然无法训练。2.5.2 Pipeline Parallelism切 LayerPipeline Parallelism 的核心是把模型按照 Layer / Stage 切开例如GPU0Layer 1 ~ 6GPU1Layer 7 ~ 12GPU2Layer 13 ~ 18GPU3Layer 19 ~ 24同一份数据依次流过GPU0 → GPU1 → GPU2 → GPU3因此 PP 主要解决完整模型太大一张 GPU 无法容纳。为了提高 GPU 利用率还需要Micro-batch让不同 Micro-batch 同时位于不同 Pipeline Stage。主要问题是Pipeline BubbleStage BalanceActivation / Gradient 通信2.5.3 Tensor Parallelism切矩阵Tensor Parallelism 切得更细。它不是把不同 Layer 分给不同 GPU而是多个 GPU 共同完成同一个 Layer例如线性层Y XW可以把W切成W1、W2、W3、W4然后GPU0 → W1GPU1 → W2GPU2 → W3GPU3 → W4多张 GPU 同时完成同一个矩阵乘法的一部分。因此 TP 主要解决单个 Transformer Layer 或大型矩阵本身就太大。它经常用于Attention HeadsQ / K / V / WoMLP 大矩阵主要问题是GPU 间通信非常频繁因此通常需要NVLink / NVSwitch等高速互联。2.5.4 DP、PP、TP 对比并行方式切分对象每张 GPU 保存什么数据如何流动主要解决的问题主要代价Data Parallelism数据完整模型不同 GPU 处理不同 Batch提高训练吞吐量梯度同步Pipeline ParallelismLayer / Stage部分模型层数据依次经过多个 Stage模型单卡放不下Pipeline BubbleTensor Parallelism单层内部矩阵矩阵的一部分多 GPU 同时算同一层单层本身太大高频 GPU 通信可以简单记成DP模型不切数据切PP模型按层切TP一层内部继续切2.5.5 三种并行方式应该怎么理解如果模型能够完整放进一张 GPU但训练太慢优先考虑 DP如果整个模型太大一张 GPU 放不下考虑 PP如果连单个 Transformer Layer 都很大考虑 TP但是超大模型训练中通常不会只选择一种。例如DP PPDP TPPP TP甚至DP PP TP一起使用。因此三种并行方式更准确的关系不是谁替代谁而是它们解决不同维度的问题可以组合使用这就自然进入下一节2.6 多种并行策略组合也就是大模型训练中常说的3D Parallelism2.6 多种并行策略组合实际的大模型训练中通常不会只使用一种并行方式。因为DP、PP、TP分别解决的是不同问题DP切数据提高训练吞吐量PP切模型层解决整个模型单卡放不下TP切单层矩阵解决单个 Layer 本身太大因此实际训练中经常组合使用DP TPDP PPTP PP以及DP PP TP其中DP PP TP通常称为3D Parallelism2.6.1 DP TP数据并行 张量并行假设有 8 张 GPU。可以划分成两个 Data Parallel GroupDP Group 0GPU0 ~ GPU3DP Group 1GPU4 ~ GPU7每个 DP Group 内部再使用 Tensor ParallelismGPU0 ~ GPU3共同完成一个模型副本中的矩阵计算。而GPU4 ~ GPU7共同完成另一份相同的模型副本。两组处理不同的数据DP Group 0 → Data ADP Group 1 → Data B因此TP负责把同一个 Layer 内部的大矩阵切开而DP负责复制多份 TP 模型处理不同数据这种组合适合模型单层较大同时还希望通过多份模型副本提高训练吞吐量。2.6.2 DP PP数据并行 流水线并行假设一个模型被划分成 4 个 Pipeline StageStage 0Stage 1Stage 2Stage 3使用 4 张 GPUGPU0 → Stage 0GPU1 → Stage 1GPU2 → Stage 2GPU3 → Stage 3这 4 张 GPU 共同组成一个完整 Pipeline如果还有另外 4 张 GPUGPU4 ~ GPU7就可以再复制一套相同的 Pipeline。于是Pipeline 0 → Data APipeline 1 → Data B两套 Pipeline 最后再进行梯度同步。因此PP负责把一个完整模型按 Layer 切开而DP负责复制多套完整 Pipeline这种组合适合整个模型太大同时还希望通过多份 Pipeline 提高整体训练吞吐量。2.6.3 TP PP张量并行 流水线并行如果模型继续增大可能出现两个问题整个模型太大同时单个 Layer 也很大这时可以同时使用TP PP先使用 Pipeline Parallelism模型 → Stage 0 Stage 1 Stage 2 Stage 3再在每个 Stage 内部使用 Tensor Parallelism。例如 8 张 GPUStage 0 → GPU0 GPU1Stage 1 → GPU2 GPU3Stage 2 → GPU4 GPU5Stage 3 → GPU6 GPU7这里PP 切不同 Layer而TP 切每个 Layer 内部的大矩阵因此可以进一步降低单张 GPU 的模型参数和计算压力2.6.4 3D ParallelismDP PP TP超大模型训练通常会同时使用Data ParallelismPipeline ParallelismTensor Parallelism这就是3D Parallelism三个维度分别负责DP → 数据维度PP → Layer 维度TP → 单层矩阵维度可以把它理解成一个训练集群↓先划分多个Data Parallel Group↓每个 DP Group 内部划分Pipeline Stage↓每个 Pipeline Stage 内部继续划分Tensor Parallel Group所以三种并行方式不是互相替代而是在不同维度共同拆分训练任务。2.6.5 GPU 数量如何计算假设DP 4PP 2TP 4那么总 GPU 数量GPU Total DP × PP × TP也就是4 × 2 × 4 32因此GPU 总数 DP × PP × TP这个公式非常重要。其中TP 4一个 Layer 内部由 4 张 GPU 共同计算PP 2一个模型被切成 2 个 Pipeline StageDP 4完整模型结构一共复制 4 份2.6.6 一个 32 GPU 的例子假设GPU Total 32设置TP 4PP 2那么一份完整模型需要TP × PP也就是4 × 2 8 张 GPU32 张 GPU 一共可以组成32 / 8 4份完整模型。所以DP 4最终DP 4PP 2TP 4并满足32 4 × 2 × 4可以理解成4 个 Data Parallel Group每个 Group 内2 个 Pipeline Stage每个 Stage 内4 张 GPU 做 Tensor Parallel2.6.7 为什么需要组合并行三种并行方式分别解决DP↓训练吞吐量不足PP↓整个模型单卡放不下TP↓单个 Layer 本身太大因此实际大模型训练通常不是DP / PP / TP 三选一而是根据模型大小和硬件条件进行组合最终目标是同时解决模型放不下训练速度太慢GPU 利用率不高的问题。2.6.8 组合并行的代价并行方式越多系统也越复杂。需要同时考虑DP 的 Gradient SynchronizationPP 的 Pipeline BubbleTP 的高频 GPU CommunicationParallel Group 的划分GPU 间通信拓扑NVLink / NVSwitch / InfiniBand显存分配计算负载均衡所以GPU 数量更多 ≠ 一定训练更快真正需要平衡的是计算显存通信并行效率因此这一部分可以最终压缩成DP 切数据PP 切 LayerTP 切矩阵进一步DP TPDP PPTP PPDP PP TP 3D Parallelism并且GPU Total DP × PP × TP这也是后面继续学习DeepSpeed和Megatron-LM的重要基础。Split Batch