3D并行与4D并行

3D并行

3D并行通常指的是将以下三种并行策略结合起来：

• D1：数据并行（Data Parallelism）
• D2：张量并行（Tensor Parallelism）
• D3：流水线并行（Pipeline Parallelism）

各个“D”的含义及详细讲解

1. 数据并行（Data Parallelism）

• 含义： 在数据并行中，每个设备（或一组设备）都持有完整的模型副本。训练数据集被分成多个批次（mini-batches），每个设备处理一个不同的批次数据。
• 工作原理：
  • 模型参数和优化器状态在所有参与训练的设备上都是复制的。
  • 每个设备接收各自的数据子集进行前向传播和反向传播，计算出梯度。
  • 在每个训练步骤结束时，所有设备上的梯度会通过**集合通信（如AllReduce）**进行聚合（通常是求平均），然后用聚合后的梯度来更新所有设备上的模型参数，确保模型参数的一致性。
• 优点： 实现相对简单，能够有效提高训练吞吐量，因为多个设备可以同时处理不同的数据。
• 缺点： 每个设备都需要存储一份完整的模型参数、梯度和优化器状态，当模型非常大时，即使一个完整的模型也可能超出单个设备的显存限制。

2. 张量并行（Tensor Parallelism，也称模型内并行或层内并行）

• 含义： 当模型太大，无法完全放入单个设备内存时，需要将模型的单个层内部的计算进行拆分，分配到多个设备上。
• 工作原理：
  • 张量并行通常针对模型中的大型矩阵乘法等操作。例如，在Transformer模型中，线性层（全连接层）或注意力机制中的权重矩阵可以被按行或按列切分，并分发到不同的设备上。
  • 每个设备只负责计算其分配到的矩阵部分。在计算完成后，需要通过**集体通信（如AllGather或AllReduce）**将部分结果聚合起来，形成完整的输出，再传递给下一层。
• 优点： 解决了单个设备内存不足以存储整个模型层参数的问题，可以在层内实现更细粒度的并行。
• 缺点： 频繁的设备间通信（特别是AllGather/AllReduce）会带来显著的通信开销，对带宽要求较高。

3. 流水线并行（Pipeline Parallelism，也称层间并行）

• 含义： 流水线并行是将模型按层（或一系列层）进行切分，将不同的模型层分配给不同的设备。
• 工作原理：
  • 模型被划分为多个连续的“阶段”（stage），每个阶段由一个或多个设备负责。
  • 在训练过程中，数据被进一步切分为更小的“微批次”（micro-batches）。
  • 一个微批次完成前向传播后，其激活值（activations）会传递给下一个阶段的设备，下一个阶段的设备可以立即开始处理这个微批次，而当前设备则可以开始处理下一个微批次，从而形成一个数据流的流水线。
  • 反向传播也以类似的方式进行，梯度沿着流水线反向传递。
• 优点： 显著减少了每个设备所需的内存，因为每个设备只存储模型的一部分层，并且通过流水线方式提高了设备的利用率，减少了空闲时间（“气泡”效应）。
• 缺点： 存在“流水线气泡”（pipeline bubble），即在流水线开始和结束时，部分设备可能处于空闲状态，降低了整体效率。微批次的大小和数量对性能影响很大。

---

4D并行

• D1：数据并行（Data Parallelism）
• D2：张量并行（Tensor Parallelism）
• D3：流水线并行（Pipeline Parallelism）
• D4：序列并行（Sequence Parallelism）

4. 序列并行（Sequence Parallelism）

• 含义： 针对处理**长序列（如长文本）**的场景，将输入序列本身进行切分，分配到不同的设备上。
• 工作原理：
  • 在Transformer等模型中，自注意力机制的计算复杂度与序列长度的平方成正比。当序列非常长时，即使激活值（activations）的内存消耗也会变得非常大。
  • 序列并行将输入序列沿着序列长度维度进行切分。每个设备只处理序列的一部分，从而显著减少每个设备上激活值的内存占用。
  • 在需要全局信息（如Attention）时，可能需要通过All-to-All通信来交换信息，但可以有效降低单卡显存压力。
• 优点： 有效解决长序列训练带来的激活值内存瓶颈。
• 缺点： 需要精心设计通信策略，可能会引入额外的通信开销。

补充 优化器状态并行（Optimizer State Parallelism / ZeRO）

• 含义： 这是一种针对 优化器状态（optimizer states）进行并行的策略。优化器状态通常包括模型参数的梯度、动量等信息，它们会占用大量的显存。
• 工作原理：
  • ZeRO (Zero Redundancy Optimizer) 系列技术（如ZeRO-1, ZeRO-2, ZeRO-3）是典型的优化器状态并行方案。
  • ZeRO-1：将优化器状态（Optimizer States）在数据并行组内的设备间进行切分。每个设备只存储优化器状态的一部分。
  • ZeRO-2：在此基础上，进一步将梯度（Gradients）也进行切分。
  • ZeRO-3：更进一步，将模型参数（Parameters）也进行切分。
• 优点： 极大减少了每个设备所需的内存，特别是对于模型参数量巨大的模型，优化器状态的存储是一个巨大的负担。
• 缺点： 引入了额外的通信，需要在前向/反向传播中动态收集和分发参数、梯度或优化器状态。

---