Skip to main content

在AI大模型训练中，单个GPU的计算能力和显存往往无法满足需求，需要使用多个GPU协同工作。并行策略就是如何合理地将训练任务分配到多个GPU上的方法。

并行策略分类对比​

并行策略核心原理简单理解适用场景
数据并行 Data Parallelism (DP)每个GPU持有完整模型，处理不同的数据分片就像4个厨师，每人拿着相同的菜谱（模型），同时做不同的菜（数据），最后交流心得（同步梯度）模型较小，数据量大
模型并行 Model Parallelism (MP)将模型的不同层分配到不同GPU就像流水线，第1个工人加工零件A（第1层），传给第2个工人加工零件B（第2层）模型超大，单GPU装不下
流水线并行 Pipeline Parallelism (PP)模型并行+微批次流水线优化改进版流水线：第1批数据在第1个工人处理时，第2批数据同时进入第2个工人大模型+想提高效率
张量并行 Tensor Parallelism (TP)将单个神经网络层的计算拆分到多个GPU就像搬一个大箱子，几个人同时抬不同部位单层计算量特别大

数据并行详解​

数据并行是最常用的并行策略，下面通过图示和术语表来深入理解：

核心术语​

术语英文定义通俗理解在数据并行中的作用
数据分片Data Sharding将完整训练数据集切分成多份，每份分配给一个GPU把一本书的页码平均分给4个人同时阅读图中GPU 0/1/2/3各自处理不同"数据分片"，避免重复计算相同数据
模型副本Model Replica每个GPU上存储的完整模型拷贝每个人手里都有一本相同的菜谱（模型）确保每个GPU都能独立完成前向和反向传播计算
前向传播Forward Pass数据从输入层经过各层计算到输出层的过程按照菜谱一步步做菜的过程每个GPU独立用自己的数据分片计算输出，互不干扰
反向传播Backward Pass从输出层计算损失，逐层计算梯度传回输入层做完菜后总结每一步的改进空间每个GPU独立计算自己数据分片产生的梯度
梯度Gradient损失函数对模型参数的导数（记为grad_W），指示参数更新方向告诉你模型参数应该往哪个方向调整、调多少每个GPU计算出"本地梯度"，需要同步后才能更新模型
梯度同步Gradient Synchronization将所有GPU的梯度汇总求平均（或求和），让所有GPU得到相同的梯度4个厨师分别做菜后，交流心得并达成一致意见数据并行的关键步骤：确保所有GPU用相同的梯度更新模型，保持模型一致性
AllReduceAllReduce一种集合通信算法，高效地完成"所有进程的数据求和并广播给所有进程"就像开会时统计大家的意见，算出平均值，再告诉每个人实现梯度同步的高效通信方式，比逐个传递快得多
参数更新Parameter Update使用梯度更新模型参数：W_new = W_old - 学习率 × grad_W根据改进建议（梯度）修改菜谱（模型）同步后所有GPU用相同梯度更新参数，确保模型一致
DistributedSamplerDistributed SamplerPyTorch提供的数据采样器，确保各GPU取到不重复的数据分片自动分配工作的调度员，保证4个人不会重复读同一页书训练时必须使用，否则所有GPU会处理相同数据导致浪费

完整训练流程（以4个GPU为例）​

  1. 初始化阶段：

    • 4个GPU各自加载相同的模型副本（模型参数完全一致）
    • 准备好完整的训练数据集

  2. 数据分片：

    • 使用DistributedSampler将一个批次（例如256张图片）均匀分成4份
    • GPU 0处理图片1-64，GPU 1处理图片65-128，以此类推

  3. 前向传播（各GPU独立计算）：

    • 每个GPU用自己的64张图片，通过模型计算得到预测结果
    • 这一步不需要通信，各GPU并行执行

  4. 反向传播（各GPU独立计算梯度）：

    • 每个GPU计算损失，并反向传播得到本地梯度
    • GPU 0得到梯度grad_W0，GPU 1得到grad_W1，GPU 2得到grad_W2，GPU 3得到grad_W3

  5. 梯度同步（关键通信步骤）：

    • 通过AllReduce操作，所有GPU交换梯度并求平均
    • 最终所有GPU得到相同的平均梯度：grad_W = (grad_W0 + grad_W1 + grad_W2 + grad_W3) / 4

  6. 参数更新（各GPU独立执行）：

    • 每个GPU用相同的平均梯度更新自己的模型参数
    • 由于起点相同（上一步参数一致）、更新量相同（同步后的梯度一致），所以更新后参数仍然一致

  7. 重复步骤2-6，直到训练完成

为什么需要梯度同步？​

假设没有梯度同步，会发生什么：

  • GPU 0看到猫的图片，学到"有毛→是猫"
  • GPU 1看到狗的图片，学到"有毛→是狗"
  • 两个GPU的模型会朝不同方向更新，导致模型"分裂"

梯度同步让所有GPU看到全局信息：

  • 同步后梯度 = "有毛→50%是猫 + 50%是狗"
  • 所有GPU用这个全局视角更新模型，保持一致性

其他并行策略​

模型并行（Model Parallelism）​

  • 原理：将模型的不同层分配到不同GPU，数据依次流过各层
  • 示例：4层神经网络，第1层在GPU 0，第2层在GPU 1...
  • 适用场景：模型太大（如GPT-3有1750亿参数），单个GPU显存装不下
  • 优势：突破单卡显存限制，可以训练超大模型
  • 缺点：同一时刻只有一个GPU在工作，其他GPU在等待，利用率低

流水线并行（Pipeline Parallelism）​

  • 原理：模型并行的优化版，将数据切分成多个小批次（微批次），流水线式处理
  • 示例：当微批次1在GPU 1处理时，微批次2可以进入GPU 0，减少空闲
  • 适用场景：大模型+需要提高GPU利用率
  • 优势：减少GPU空闲时间，提高资源利用率
  • 实现：如DeepSpeed、Megatron-LM等框架

流水线示意：

时间 t1: [GPU0处理batch1-layer1] [GPU1空闲]         [GPU2空闲]         [GPU3空闲]
时间 t2: [GPU0处理batch2-layer1] [GPU1处理batch1-layer2] [GPU2空闲]         [GPU3空闲]
时间 t3: [GPU0处理batch3-layer1] [GPU1处理batch2-layer2] [GPU2处理batch1-layer3] [GPU3空闲]
时间 t4: [GPU0处理batch4-layer1] [GPU1处理batch3-layer2] [GPU2处理batch2-layer3] [GPU3处理batch1-layer4]

张量并行（Tensor Parallelism）​

  • 原理：将单个网络层的矩阵运算拆分到多个GPU并行计算
  • 示例：一个[1024, 1024]的矩阵乘法，拆成4个[1024, 256]的小矩阵在4个GPU上同时计算
  • 适用场景：超大的Transformer模型，单个注意力层计算量巨大
  • 优势：层内并行，减少等待时间
  • 实现：如Megatron-LM对GPT模型的张量并行实现

矩阵分割示例：

原始矩阵乘法: [M × K] × [K × N] = [M × N]

张量并行（按列切分）:
GPU 0: [M × K] × [K × N/4] = [M × N/4]
GPU 1: [M × K] × [K × N/4] = [M × N/4]
GPU 2: [M × K] × [K × N/4] = [M × N/4]
GPU 3: [M × K] × [K × N/4] = [M × N/4]

最后拼接: [M × N/4] + [M × N/4] + [M × N/4] + [M × N/4] = [M × N]

混合并行策略​

在实际的大模型训练中，通常会组合使用多种并行策略以获得最佳性能：

2维并行：数据并行 + 模型并行​

适用于中大型模型训练：

  • 场景：8节点，每节点8卡，共64卡
  • 策略：节点间数据并行（8路），节点内模型并行（8层）
  • 优势：既能利用数据并行的高效性，又能突破单卡显存限制

3维并行：数据并行 + 张量并行 + 流水线并行​

适用于超大模型训练（如GPT-3）：

  • 数据并行：跨多个节点复制模型
  • 流水线并行：将模型垂直切分为多个stage
  • 张量并行：在每个stage内，将层水平切分

示例配置（GPT-3训练）：

  • 数据并行度：8
  • 流水线并行度：16
  • 张量并行度：8
  • 总GPU数：8 × 16 × 8 = 1024

通信开销对比​

并行策略通信频率通信数据量带宽需求延迟敏感度
数据并行每个batch梯度大小 ≈ 模型大小高中
模型并行每层前向/反向激活值大小中高
流水线并行每个微批次激活值大小低中
张量并行每层内部中间结果极高极高

参考资料​