跳转至

并行策略

学习目标

  1. 理解数据并行、模型并行、流水线并行的原理和适用场景
  2. 掌握各种并行策略的优势和局限性
  3. 了解混合并行框架Megatron-LM和DeepSpeed的设计要点
  4. 熟悉通信优化中的梯度压缩与计算通信重叠技术
  5. 能够根据实际场景选择合适的并行策略

3.1 并行策略概述

3.1.1 为什么需要并行策略

训练大规模神经网络模型需要处理海量的计算和存储任务。以参数量为1750亿的GPT-3为例，仅模型参数就需要约350GB的存储空间（FP16），加上优化器状态、梯度和激活值，单个样本可能需要数GB的显存。显然，单张GPU无法容纳这样的工作负载。

并行策略是解决这一问题的核心技术。通过将模型和数据分散到多个计算设备上，并行执行计算任务，我们可以：

  1. 突破显存限制：使训练超大模型成为可能
  2. 加速训练过程：多个设备同时工作，缩短训练时间
  3. 提高资源利用率：充分利用集群中的计算资源

3.1.2 并行策略的分类

根据分割方式的不同，分布式并行策略可以分为三大类：

数据并行（Data Parallelism, DP）：

  • 将训练数据分成多个子集，分配到不同设备上
  • 每个设备持有完整的模型副本
  • 适合数据量大、模型相对较小的场景
  • 是目前最广泛使用的并行策略

模型并行（Model Parallelism, MP）：

  • 将模型分割到多个设备上
  • 每个设备只负责模型的一部分计算
  • 适合模型巨大、无法放入单卡显存的场景
  • 包括张量并行和流水线并行

混合并行（Hybrid Parallelism）：

  • 结合多种并行策略
  • 在不同维度上分割模型和数据
  • 适合超大规模训练场景
┌────────────────────────────────────────────────────────────┐
│                      并行策略分类                          │
├────────────────────────────────────────────────────────────┤
│                                                            │
│   ┌──────────────────┐      ┌──────────────────┐         │
│   │    数据并行      │      │    模型并行       │         │
│   │                  │      │                  │         │
│   │  ┌────────────┐ │      │  ┌───┐ ┌───┐ ┌───┐ │         │
│   │  │ GPU0: Model│ │      │  │GPU│ │GPU│ │GPU│ │         │
│   │  │  + Data1  │ │      │  │ 0 │ │ 1 │ │ 2 │ │         │
│   │  └────────────┘ │      │  │P0 │ │P1 │ │P2 │ │         │
│   │  ┌────────────┐ │      │  └───┘ └───┘ └───┘ │         │
│   │  │ GPU1: Model│ │      │   张量并行          │         │
│   │  │  + Data2  │ │      ├──────────────────────┤         │
│   │  └────────────┘ │      │  ┌───┐ ┌───┐ ┌───┐ │         │
│   │  ┌────────────┐ │      │  │GPU│ │GPU│ │GPU│ │         │
│   │  │ GPU2: Model│ │      │  │ 0 │ │ 1 │ │ 2 │ │         │
│   │  │  + Data3  │ │      │  │L0-│ │L2-│ │L4-│ │         │
│   │  └────────────┘ │      │  │L1 │ │L3 │ │L5 │ │         │
│   └──────────────────┘      │  └───┘ └───┘ └───┘ │         │
│                             │   流水线并行          │         │
│                             └──────────────────────┘         │
└────────────────────────────────────────────────────────────┘

3.1.3 选择并行策略的考量因素

选择并行策略需要综合考虑多种因素：

模型规模：

  • 百亿参数以下：数据并行通常足够
  • 百亿到万亿参数：需要模型并行或混合并行
  • 万亿参数以上：必须使用混合并行

集群配置：

  • GPU数量和互联带宽
  • 节点内和跨节点带宽差异
  • 存储系统性能

通信开销：

  • 梯度同步频率和数据量
  • 集合通信的效率
  • 网络拓扑的影响

实现复杂度：

  • 不同并行策略的开发难度
  • 现有框架的支持程度
  • 调试和维护成本

3.2 数据并行

3.2.1 数据并行的原理

数据并行是最直观且应用最广泛的并行策略。其核心思想是：在每个计算设备上保留完整的模型副本，但处理不同的数据子集。

数据并行的工作流程：

  1. 数据分片：将训练数据分成多个mini-batch，每个设备处理一份
  2. 并行前向传播：各设备使用相同的模型参数，但输入不同的数据
  3. 梯度计算：各设备独立计算梯度
  4. 梯度同步：通过AllReduce操作汇总所有设备的梯度
  5. 参数更新：使用汇总后的梯度更新模型参数

关键假设：

数据并行基于这样一个假设：各设备上的模型参数相同，因此可以独立计算梯度，之后通过同步确保参数一致。

3.2.2 同步数据并行与异步数据并行

根据梯度同步方式的不同，数据并行可以分为：

同步数据并行（Synchronous Data Parallelism）：

  • 所有设备在每一轮迭代后同步梯度
  • 参数更新等待所有设备完成计算
  • 训练过程与单设备等价，收敛性有保证

优点：收敛行为与单设备训练一致 缺点：需要等待最慢的设备，存在"木桶效应"

异步数据并行（Asynchronous Data Parallelism）：

  • 各设备独立进行计算和参数更新
  • 不等待其他设备，直接使用最新参数
  • 可能导致"梯度过期"问题

优点：设备利用率高，无等待时间 缺点：收敛性可能受影响，需要特殊处理

实际应用中，同步数据并行更为常见，因为其收敛行为更容易预测。

3.2.3 分布式数据并行（DDP）

分布式数据并行是PyTorch等框架中实现数据并行的标准方式，相比早期的DataParallel（DDP仅支持单节点），DDP支持多节点扩展且效率更高。

DDP的核心设计：

  1. 多进程架构：每个GPU一个进程，绕过了Python GIL限制
  2. Ring AllReduce：高效的梯度同步算法
  3. 延迟隐藏：梯度计算与通信重叠，减少等待时间

Ring AllReduce算法：

Ring AllReduce将梯度同步分成两个阶段：

阶段一：Reduce-Scatter - 每个节点将自己的梯度分成N份 - 进行N-1轮传递，每轮累加一份数据 - 最终每份数据汇聚了所有节点的结果

阶段二：AllGather - 同样进行N-1轮传递 - 每轮交换数据，最终每节点获得完整汇总结果

Ring AllReduce的优势是通信量与节点数无关，适合大规模集群。

3.2.4 DDP的实现细节

PyTorch DDP的实现包含几个关键组件：

梯度桶（Gradient Bucketing）：

  • 不是为每个参数单独触发AllReduce
  • 而是将多个参数的梯度放入同一个桶
  • 桶满后触发一次集合通信
  • 减少通信次数，提高带宽利用率

延迟反向传播（Backward Overlap）：

  • 在反向传播进行时就开始通信已计算好的梯度
  • 通信与计算并行执行
  • 有效隐藏通信延迟

同步检查点（Sync BatchNorm）：

  • BatchNorm层需要统计全局数据
  • DDP在训练前同步各设备的BatchNorm统计量
  • 确保训练行为一致

3.2.5 完全分片数据并行（FSDP）

当模型规模进一步增大时，单纯的DDP会遇到显存瓶颈。完全分片数据并行（Fully Sharded Data Parallel, FSDP）通过在多个设备间分片模型状态来大幅降低显存需求。

ZeRO优化器：

FSDP基于微软提出的ZeRO（Zero Redundancy Optimizer）技术，包括三个阶段：

阶段 分片内容 显存节省
ZeRO-1 优化器状态 4倍
ZeRO-2 优化器状态 + 梯度 8倍
ZeRO-3 优化器状态 + 梯度 + 参数 与设备数成正比

ZeRO-1的工作原理：

  1. 将优化器状态（如Adam的momentum和variance）分片到各设备
  2. 前向/反向传播在各设备独立进行
  3. 梯度通过Reduce-Scatter聚合，每个设备只更新其分片的参数
  4. 参数更新后通过AllGather收集完整参数

ZeRO-3的通信模式：

前向传播时，需要通过AllGather收集跨设备的参数分片；反向传播完成后，同样需要聚合梯度。这种以通信换显存的思想，使超大模型训练成为可能。

3.2.6 数据并行的挑战与局限

尽管数据并行是应用最广泛的并行策略，但它也存在一些局限性：

显存瓶颈：

  • 每个设备都需要保存完整的模型副本
  • 当模型规模超过单卡显存时，数据并行无法直接使用
  • 需要结合模型并行来解决

通信开销：

  • 每轮迭代都需要同步梯度
  • 当GPU数量增加时，通信时间可能成为瓶颈
  • 需要通信优化技术来缓解

扩展效率：

  • 随着GPU数量增加，扩展效率会下降
  • 主要是通信开销增加导致的
  • 大规模训练需要仔细调优

3.3 模型并行

3.3.1 模型并行的概念

当模型规模超过单卡显存时，数据并行无法直接使用。模型并行（Model Parallelism）通过将模型分割到多个设备上来解决这个问题。

模型并行的核心思想是：不同设备负责模型的不同部分，可以是不同层，也可以是同一层的不同神经元。

模型并行的类型：

  1. 张量并行（Tensor Parallelism, TP）：将单层内部的操作分割到多个设备
  2. 流水线并行（Pipeline Parallelism, PP）：将不同层分配到不同设备

3.3.2 张量并行

张量并行将单层计算（如矩阵乘法）的操作数分割到多个设备上，实现单层的并行计算。

矩阵乘法的张量并行：

以矩阵乘法 \(Y = X \cdot W\) 为例： - 将 \(W\) 按列分割成 \([W_1, W_2]\)，分别放在GPU0和GPU1上 - \(X\) 同时发送到两个GPU - 各GPU计算部分结果 \(Y_1 = X \cdot W_1\), \(Y_2 = X \cdot W_2\) - 通过AllGather汇总完整结果

Transformer中的张量并行：

Transformer的主要计算是矩阵乘法和注意力机制。

对于MLP层： - 第一个线性层的权重矩阵 \(W_1\) 按列分割 - 第二个线性层的权重矩阵 \(W_2\) 按行分割 - 中间结果通过AllReduce同步

对于注意力层： - Q、K、V矩阵按列分割 - 输出投影 \(W_O\) 按行分割 - 各GPU计算注意力的部分结果后汇总

3.3.3 张量并行的通信模式

张量并行的通信发生在每层计算之后，需要仔细考虑通信开销。

通信模式分析：

  • 张量并行通信量与层输出维度成正比
  • 每层前向传播和反向传播各需要一次AllReduce
  • 对于Transformer，约占总时间的15-20%

通信优化策略：

  1. 算子融合：将多个小通信操作合并
  2. 异步通信：计算与通信重叠
  3. 拓扑感知：优先将张量并行限制在节点内

3.3.4 朴素模型并行与流水线并行

在张量并行之外，还有两种模型并行的经典方式：

朴素模型并行（Naive Model Parallelism）：

最直观的模型分割方式：将模型的不同层分配到不同设备。

例如，12层Transformer模型分到4个GPU： - GPU0：层0-2 - GPU1：层3-5 - GPU2：层6-8 - GPU3：层9-11

问题：大部分时间只有一个GPU在工作，其他GPU空闲。

流水线并行（Pipeline Parallelism）：

流水线并行通过将输入数据分成多个微批次（micro-batch），让不同设备处理不同批次的工作。

工作流程： - GPU0处理micro-batch 1的前向传播 - GPU0处理micro-batch 1时，GPU1可以同时处理micro-batch 0的反向传播 - 形成类似工业流水线的生产过程

3.3.5 Gpipe流水线并行

Gpipe是Google提出的流水线并行方案，通过微批次（micro-batch）和重计算（Recomputation）来实现高效的流水线。

微批次处理：

Gpipe将mini-batch进一步分成更小的micro-batch： - 设mini-batch大小为 \(m\)，分成 \(n\) 个micro-batch - 每个micro-batch的处理时间更短，便于流水线调度

F-then-B策略：

Gpipe采用先前向（F）后反向（B）的调度策略： - 先完成所有micro-batch的前向传播 - 再开始所有micro-batch的反向传播 - 简化了实现，但需要保存所有激活值

重计算策略：

为了减少显存占用，Gpipe在前向传播时不保存激活值，只保存每个micro-batch的输入和输出： - 在反向传播时，重新计算前向传播 - 以计算换显存，激活值显存减少约50%

3.3.6 PipeDream流水线并行

PipeDream是微软提出的另一种流水线并行方案，采用1F1B（One Forward One Backward）策略。

1F1B策略：

与Gpipe不同，PipeDream在完成一个micro-batch的前向传播后，立即执行其反向传播：

  • 前向传播完成后立即反向传播
  • 释放该micro-batch相关的激活值
  • 大幅减少显存占用

权重版本管理：

由于流水线中不同阶段使用的参数版本不同，PipeDream需要管理权重版本：

  • 每个阶段维护自己的参数副本
  • 前向传播使用当前参数
  • 反向传播使用对应的参数版本
  • 参数更新后通知下游阶段

Round-Robin调度：

PipeDream还支持将同一阶段的数据并行任务分配到多个GPU，通过轮询调度保证每个micro-batch的前向和反向在同一GPU上完成。

3.4 混合并行

3.4.1 混合并行的需求

当模型规模达到万亿参数时，单一的并行策略往往无法满足需求：

  • 数据并行：需要完整模型副本，显存不足
  • 流水线并行：流水线气泡影响效率
  • 张量并行：通信开销随并行度增加

混合并行结合多种策略，在不同维度上分割模型和数据，最大化利用集群资源。

3.4.2 3D混合并行

3D混合并行是当前训练超大规模模型的主流方案，由数据并行（DP）、流水线并行（PP）和张量并行（TP）三个维度组成。

3D并行的组织：

以8个GPU为例，2×2×2的3D并行： - TP维度=2：张量并行，两GPU共享一层计算 - PP维度=2：流水线并行，两GPU负责不同层的子集 - DP维度=2：数据并行，两GPU处理不同数据

内存效率：

  • TP减少每层的显存占用
  • PP减少总的层数
  • DP通过ZeRO进一步分片优化器状态

计算效率：

  • TP优先放在节点内，利用NVLink高带宽
  • PP跨节点开销最小，因为通信量最少
  • DP的梯度同步可以与计算重叠

3.4.3 Megatron-LM

Megatron-LM是NVIDIA开发的大规模 transformer训练框架，专为张量并行和混合并行优化。

核心特性：

  1. 高效张量并行：实现了Multi-Head Attention和MLP的高效张量并行
  2. 序列并行：在序列维度上分割LayerNorm等操作，节省激活显存
  3. 选择性重计算：只重计算注意力等耗时操作
  4. FlashAttention集成：加速注意力计算

张量并行实现：

Megatron-LM的张量并行实现包含：

# 简化示例
class TensorParallelLinear(nn.Module):
    def __init__(self, config, partition_dim=0):
        self.weight = partition_weights(config, partition_dim)

    def forward(self, x):
        # 本地计算
        local_output = F.linear(x, self.weight)
        # 集合通信
        output = all_gather(local_output, dim=-1)
        return output

与数据并行结合：

Megatron-LM配合张量并行使用，需要注意： - 张量并行需要在节点内进行（NVLink带宽） - 数据并行跨节点进行 - 中间结果需要路由到正确的GPU

3.4.4 DeepSpeed

DeepSpeed是微软开发的大规模训练优化框架，专注于提高训练效率和降低显存需求。

ZeRO优化：

DeepSpeed实现了完整的ZeRO优化器： - ZeRO-1：优化器状态分片 - ZeRO-2：优化器状态+梯度分片 - ZeRO-3：全量参数分片 - ZeRO-Infinity：将分片扩展到CPU和NVMe内存

3D并行支持：

DeepSpeed支持完整的3D并行配置：

# DeepSpeed配置文件示例
{
    "train_batch_size": 1024,
    "zero_optimization": {
        "stage": 3,
        "cpu_offload": true,
        "offload_param": {
            "pin_memory": true
        }
    },
    "gradient_clipping": 1.0,
    "prescale_gradients": false,
    "steps_for_print_qeury": 10
}

通信优化：

DeepSpeed实现了多种通信优化： - 梯度压缩（PowerSGD等） - 通信与计算重叠 - 异步数据加载

3.4.5 ColossalAI

ColossalAI是潞晨科技开发的大规模训练平台，集成了多种并行策略。

核心特性：

  1. 多维并行：原生支持DP、PP、TP及其组合
  2. 异构训练：支持将部分张量卸载到CPU内存
  3. 便捷接口：提供简洁的并行化API
  4. 自动并行：自动搜索最优并行策略

集成策略：

ColossalAI的集成方式允许灵活组合： - 可以只使用某种并行策略 - 也可以开启完整的3D并行 - 配置简单，与现有代码改动最小

3.5 通信优化

3.5.1 通信优化的重要性

在分布式训练中，通信开销往往是限制扩展效率的主要因素。通信优化的目标是减少通信数据量、降低通信延迟、提高带宽利用率。

通信开销的来源：

  1. 梯度同步：数据并行中，每轮迭代需要同步梯度
  2. 参数同步：模型并行中，不同阶段需要传递中间结果
  3. 数据加载：多个节点同时读取训练数据

优化方向：

  1. 减少通信数据量（梯度压缩、量化）
  2. 隐藏通信延迟（计算通信重叠）
  3. 提高通信效率（拓扑感知路由）

3.5.2 梯度压缩

梯度压缩通过减少传输的梯度数据量来加速通信。

量化压缩（Gradient Quantization）：

将32位浮点梯度压缩为低精度表示：

格式 精度 压缩率
FP16 16位 2x
INT8 8位 4x
INT4 4位 8x

量化压缩的挑战： - 低精度可能导致梯度精度损失 - 需要选择合适的量化策略（如随机舍入）

Top-K稀疏压缩：

只传输梯度中最重要的k个元素：

  1. 计算每个梯度的绝对值或动量
  2. 选择绝对值最大的k个元素
  3. 传输这k个元素的值和索引
  4. 接收端重建稀疏梯度

压缩率 = 梯度总数 / k，稀疏度越高压缩率越高，但可能影响收敛。

随机压缩（Random Sparsification）：

与Top-K不同，随机压缩按概率选择元素： - 每个元素以概率p传输 - 接收端通过期望值估计 - 适合通信不稳定的环境

3.5.3 计算与通信重叠

计算与通信重叠是隐藏延迟的核心技术。

重叠的基本原理：

  1. 将训练过程分成多个阶段
  2. 在计算当前阶段的同时，异步执行前一阶段的通信
  3. 两阶段完成后再进入下一阶段

PyTorch DDP的重叠实现：

DDP使用梯度桶和钩子函数实现重叠：

# 伪代码示例
for batch in dataloader:
    # 启动多个bucket的梯度计算
    for bucket in buckets:
        start_backward(bucket.grads)

        # 梯度就绪后触发AllReduce
        if bucket.ready():
            async_allreduce(bucket)

重叠的条件：

并非所有情况下都能实现重叠： - 需要足够的算力支撑通信期间仍有任务执行 - 通信操作不能阻塞后续计算 - 需要仔细的调度策略

3.5.4 通信原语优化

针对集合通信原语的优化也是重要方向。

AllReduce算法选择：

不同的AllReduce算法适用于不同场景： - Ring AllReduce：适合高带宽、节点数较多的场景 - Tree AllReduce：适合低延迟、节点数较少的场景 - Rabenseifner's算法：基于拓扑的优化算法

拓扑感知通信：

利用网络拓扑信息优化通信路径： - 节点内通信使用NVLink（最高900 GB/s） - 跨节点通信使用InfiniBand（200 Gbps） - 优先在节点内完成更多通信

NCCL自动优化：

NCCL库能够自动检测网络拓扑并选择最优通信算法： - 检测GPU之间的连接方式 - 选择最合适的Ring/Tree配置 - 自动调整缓冲区大小

3.5.5 自适应通信策略

根据训练状态动态调整通信策略。

基于带宽的自适应：

不同训练阶段对带宽需求不同： - 训练初期：参数变化大，需要高精度同步 - 训练后期：参数变化小，可以降低同步精度 - 评估阶段：可以跳过某些同步

基于延迟的自适应：

网络状况可能动态变化： - 带宽下降时增加压缩率 - 延迟增加时调整通信策略 - 故障时切换到备选路径

分层通信：

将通信分成多个层级： - 层内通信：使用高精度、本地链路 - 层间通信：使用较低精度、跨节点链路 - 根据数据重要性选择通信层级

小结

本节我们系统学习了分布式训练中的并行策略和通信优化技术。

核心要点：

  1. 数据并行通过在多个设备上复制完整模型、处理不同数据子集来加速训练，是目前应用最广泛的并行策略
  2. 分布式数据并行（DDP）使用Ring AllReduce算法实现高效的梯度同步，是现代分布式训练的基础
  3. 完全分片数据并行（FSDP）基于ZeRO技术，通过分片模型状态来大幅降低显存需求
  4. 模型并行通过将模型分割到多个设备来解决单卡显存不足的问题
  5. 张量并行将单层计算分割到多个设备，适合层内部存在大量计算的场景
  6. 流水线并行将不同层分配到不同设备，通过micro-batch实现计算与通信的重叠
  7. 混合并行结合数据并行、流水线并行和张量并行，适合万亿参数规模的训练
  8. Megatron-LM和DeepSpeed是当前主流的混合并行训练框架
  9. 通信优化通过梯度压缩、计算通信重叠等技术提升分布式训练效率

思考题：

  1. 数据并行在什么情况下会遇到瓶颈？FSDP是如何解决这个问题的？
  2. Ring AllReduce算法为什么比简单的树形AllReduce更适合大规模集群？
  3. 流水线并行中的"气泡"是如何产生的？有哪些方法可以减少气泡？
  4. 3D混合并行中，为什么要将张量并行限制在节点内？
  5. 梯度压缩技术（如Top-K稀疏）在实际应用中需要注意哪些问题？