Accelerate ND-Parallel：高效多 GPU 训练指南

由于不同并行策略的复杂性，在多个 GPU 上训练大型模型可能充满挑战。在 Accelerate 中，我们与 Axolotl 合作，集成了一种快速简便的方法，可以在您的训练脚本中使用任何并行策略组合！

以下是如何将其添加到您的训练脚本中

from transformers import AutoModelForCausalLM
from accelerate import Accelerator
from accelerate.parallelism_config import ParallelismConfig
from accelerate.utils import FullyShardedDataParallelPlugin

# configure your desired parallelisms here - this particular configuration requires at least 2 nodes with 8 GPUs each. 
# setting any parallelism degree to 1 disables it i.e. dp_replicate_size=1 disables DP.
pc = ParallelismConfig(
    dp_shard_size=2, # Fully Sharded Data Parallel degree
    dp_replicate_size=2, # Data Parallel degree
    cp_size=2, # Context Parallel degree
    tp_size=2, # Tensor Parallel degree
)

fsdp_plugin = FullyShardedDataParallelPlugin(
    fsdp_version=2,
    auto_wrap_policy="transformer_based_wrap",
    transformer_cls_names_to_wrap=["LlamaDecoderLayer"],
    state_dict_type="SHARDED_STATE_DICT",
)

accelerator = Accelerator(
    parallelism_config=pc,
    fsdp_plugin=fsdp_plugin
)

model = AutoModelForCausalLM.from_pretrained(
    "NousResearch/Hermes-3-Llama-3.1-8B", 
    device_mesh=accelerator.torch_device_mesh
)

model = accelerator.prepare(model)

我们还在 Accelerate 存储库中包含了一个更全面的端到端训练脚本，其中演示了如何设置数据加载器、优化器和训练循环，以及如何在训练后保存模型。

为了进一步简化大规模模型微调并结合并行策略与各种微调技术，我们还将此技术集成到 Axolotl 中。为了帮助您立即上手，我们测试了一些示例配置，您可以根据自己的需求进行修改 - 尝试使用以下命令：

# note: this requires a minimum world size of 16 
axolotl train examples/distributed-parallel/llama-3_1-8b-hsdp-tp.yaml

您还可以查看 Axolotl ND-Parallelism 文档以获取更多详细信息——将 ND 并行技术添加到您现有配置中就像在您的 Axolotl 配置文件中添加一个或多个以下字段一样简单：

# Fully Sharded Data Parallel degree (note: also requires the fsdp_config field) 
# see https://docs.axolotl.ai/docs/multi-gpu.html#sec-fsdp for more details
dp_shard_size: 2
# Data Parallel degree
dp_replicate_size: 2
# Context Parallel Degree
context_parallel_size: 2
# Tensor Parallel Degree
tensor_parallel_size: 2

我们已经通过 Accelerate 中的 ParallelismConfig 类或 Axolotl 中的配置字段，使得配置不同并行策略的程度以及它们如何组合变得容易。但是我们如何知道哪种配置最适合我们的用例呢？当我们扩展到训练具有数百亿甚至数千亿参数的模型时，主要的挑战来自于理解不同的并行策略以及它们如何相互作用以最小化设备间的通信开销。在这篇文章中，我们将详细介绍不同的并行策略如何工作，以及何时以及如何组合它们。

目录

• 数据并行
• 完全分片数据并行
• 张量并行
• 上下文并行
• ND 并行
  • 混合分片数据并行
  • 完全分片数据并行 + 张量并行
  • 完全分片数据并行 + 上下文并行
  • 混合分片数据并行 + 张量并行
• 使用注意事项

数据并行

数据并行 (DP) 是在多个 GPU 上训练模型最常用的技术，它涉及在每个设备上复制模型、梯度和优化器状态，同时在 GPU 之间均匀分配数据批次，并在更新参数之前同步设备间的梯度。与单设备训练相比，这可以显著提高吞吐量，但要求您的模型能够适应单个设备。

我们可以通过 Accelerate 的 `ParallelismConfig` 中的 `dp_replicate_size` 参数或 Axolotl 中的配置字段来控制模型的副本数量。值得注意的是，DP 是一种*最顶层*的并行策略，这意味着如果我们使用 `dp_replicate_size=2` 并将其与其他并行策略组合，将会有 2 个模型副本，每个副本也会受到其他并行策略的影响。例如，如果我们将 `dp_replicate_size=2` 和 `tp_size=2` 结合使用，我们将拥有 2 个模型副本，每个副本都有 2 个张量并行分片。

我们使用术语 *shard* 来描述单个设备上的数据，它是较大数据的分区。

完全分片数据并行

如果我们的模型太大，无法适应单个设备，该怎么办？完全分片数据并行 (FSDP) 通过在 GPU 之间分片（均匀分布）模型的权重、梯度和优化器状态来解决此问题（这受 DeepSpeed 的 ZeRO-3 启发），同时每个设备仍接收其完整数据批次的一部分。如您从上图中所注意到的，我们不是在每个设备上都需要整个模型的完整副本，而是在前向传播之前一次只收集一个层的权重，之后权重可以再次分片。

通过这种方式，我们以内存使用量换取了在每次前向和后向传播之前收集分片参数以及进行 reduce-scatter 本地梯度的通信开销。我们可以通过调整参数收集的粒度来控制 FSDP 中的这种权衡。在一个极端情况下，我们可以对模型的每一层进行收集和重新分片，这将导致最低的峰值内存使用量，但会产生最高的通信成本。在实践中，一种常见的方法是一次性收集整个 Transformer 解码器块的权重。

虽然我们可以进一步进行内存-计算的权衡，并将模型参数和梯度卸载到 CPU 以训练更大的模型，但这可能会非常慢。相反，让我们考虑如何有效利用更多设备来训练更大的模型，同时保持高数据吞吐量。

我们使用术语 *节点* 指代托管多个 GPU（最多 8 个）的单台机器，其中 GPU 之间使用 NVLink 等实现快速节点内通信。在多节点训练中，我们依靠 Infiniband 等相对较慢的节点间通信通道。我们还将进程池中的设备总数称为世界大小——例如，一台拥有 8 个 GPU 的单节点表示世界大小为 8，而 4 个节点则表示世界大小为 32。

当在多个节点上使用 FSDP 时，我们将跨节点的所有设备视为在单个节点上进行训练。例如，对于 4 个节点，每个节点包含 8 个 GPU，我们跨 32 个设备执行分片，并使用节点内和节点间通信后端执行集体 all-reduce 和 reduce-scatter 操作。通过这种方式，FSDP 单独就可以扩展到大量 GPU，并具有较大的全局批量大小以提高数据吞吐量。然而，在某些情况下会出现一些挑战，可能需要将 FSDP 与其他并行技术结合使用。我们通常会尽量避免在超过一个完整节点的情况下使用 FSDP，因为通信开销可能会变得过高，我们将在混合分片数据并行部分讨论如何解决这个问题。

您可以使用 Accelerate 的 `ParallelismConfig` 中的 `dp_shard_size` 参数，结合已准备好的 FullyShardedDataParallelPlugin，或者在 Axolotl 中设置 `dp_shard_size` 配置字段来设置应用于模型的 FSDP 程度。

张量并行

张量并行 (TP) 是一种模型并行技术，其中模型的分片永久存储在不同的设备上，与数据并行技术相反，每个设备接收相同批次的数据。TP 通过在设备之间分配线性层的计算来工作，因此每个设备只计算矩阵乘法的一部分。这种技术最适用于大型线性层，例如 transformer 模型中的前馈层，这些层可以跨设备进行拆分。我们还可以在注意力层中的每个查询、键、值和输出投影上使用 TP，几乎没有额外的通信成本。

为了达到最佳性能，连续层的参数可以以特定方式分布，从而最大限度地减少所需的通信。当处理成对的线性层时，我们可以对第一层进行列式拆分，对后续层进行行式拆分，从而只需一次 all-reduce 操作即可组合分片输出。

与 FSDP 的动态分片行为不同，TP 创建静态内存分区，从而导致内存使用量随着 TP 组大小的增加而恒定减少。这对于大型模型至关重要，因为即使是单个解码器层也太大，无法在 FSDP all-gather 期间放入内存（回想一下 FSDP 的常见做法是同时收集整个解码器层的权重）。然而，与 FSDP 在节点间相对线性地扩展（在同构集群上最多可达约 512 个 GPU，在低带宽连接上则显著减少）不同，TP 仅在一个节点范围内有效。TP 在计算过程中需要设备之间频繁进行激活同步，因为每个设备只计算输出的一部分，需要与其他设备的输出进行通信才能继续前向传播。因此，如果要在多节点设置中使用 TP，我们必须考虑将 TP 与其他并行技术结合使用，同时将 TP 仅限于单个节点。由于其巨大的通信开销，不建议将 TP 用于 PCIe 连接的 GPU。

在 Accelerate 中，TP 大小通过 `ParallelismConfig` 中的 `tp_size` 进行配置，而在 Axolotl 中，您可以使用 `tensor_parallel_size` 配置字段。

上下文并行

最近，大型语言模型 (LLM) 的推理能力导致序列长度急剧增加，因为模型使用越来越多的 token 来解决复杂任务。为了通过微调实现这种行为，我们需要一种方法来训练模型处理非常长的序列长度——有时甚至可以达到一百万个 token！

由于 transformer 中的注意力操作与上下文长度呈平方关系，这使得在单个 GPU 上进行操作变得不可能。例如，在微调相对较小的模型（如 Mistral-7B，使用 32 个注意力头）时，如果序列长度为 128k，单个注意力矩阵将占用 128k * 128k * 2 字节 * `num_heads=32` = ~32GB * 32 = ~1TB 的激活内存！尽管在使用 FlashAttention 等优化注意力实现时这个例子不现实，但它有助于说明上下文长度增加所导致的内存需求增长。

通过上下文并行 (CP)，我们可以沿序列维度对输入进行分片，从而使每个设备只处理完整上下文的一部分，并计算完整且非常大的注意力矩阵的较小部分。为了了解其工作原理，请回忆注意力计算由以下方程描述： $\text{Attention}(Q, K, V) = \text{softmax}(QK^T)V$

其中 $Q$、$K$ 和 $V$ 分别是查询、键和值矩阵。$Q$ 的每个查询向量（行或输入嵌入）必须计算与整个序列中 $K$ 的*每个*键向量的注意力得分，以正确应用 softmax 归一化。然后，这些注意力得分将与 $V$ 中的*所有*值向量进行加权。

这里最关键的细节在于，$Q$ 中的每一行都可以独立计算其注意力分数，但每个查询向量仍然需要完整的 $K$ 和 $V$ 矩阵。换句话说，给定一个序列长度为 $n$ 的输入，我们可以将上述注意力方程扩展为

$$\begin{align} \text{Attention}(Q, K, V)_1 &= \text{softmax}(Q_1 K^T) V \\ \text{Attention}(Q, K, V)_2 &= \text{softmax}(Q_2 K^T) V \\ &\vdots \\ \text{Attention}(Q, K, V)_n &= \text{softmax}(Q_n K^T) V \end{align}$$

其中我们把查询矩阵的每一行表示为 $Q_1, Q_2, ..., Q_n$。这可以推广为：$\text{Attention}(Q, K, V)_i = \text{softmax}(Q_i K^T) V \quad \forall i \in \{1, 2, ..., n\}$

当我们跨设备对输入进行分片时，由此产生的 $Q$、$K$ 和 $V$ 矩阵（由这些输入分片计算得出）也会沿序列维度自动分片——每个 GPU 仅为其序列部分计算查询、键和值。例如，如果世界大小为 $W$ 个 GPU，序列长度为 $n$

• GPU 0 计算 $Q_{1:n/W}$、$K_{1:n/W}$、$V_{1:n/W}$
• GPU 1 计算 $Q_{n/W+1:2n/W}$、$K_{n/W+1:2n/W}$、$V_{n/W+1:2n/W}$
• ...
• GPU $(W-1)$ 计算 $Q_{(W-1)n/W+1:n}$、$K_{(W-1)n/W+1:n}$、$V_{(W-1)n/W+1:n}$

我们如何确保注意力计算正确？如上所述，每个设备只需要自己的 $Q$ 分片，但需要完整的 $K$ 和 $V$ 矩阵才能正确计算注意力。我们可以通过使用一种称为 环注意力（RingAttention）的技术来实现这一点，其工作原理如下：

1. 最初，每个 GPU 都拥有其分片的 $Q$、$K$、$V$ (例如，GPU 0 拥有 $Q_{1:n/W}$、$K_{1:n/W}$、$V_{1:n/W}$)。
2. 每个 GPU 然后为其 $Q_i$ 分片及其本地 $K_j$ 和 $V_j$ 分片计算一个部分注意力矩阵 $A_{i,j}$。
3. 每个 GPU 将其 $K$、$V$ 分片发送到环中的下一个 GPU。
4. 每个 GPU 都会从环中的上一个 GPU 接收到不同的 K 和 V 分片。
5. 每个 GPU 使用接收到的 $K$、$V$ 分片计算额外的部分注意力矩阵 $A_{i,j+1}$、$A_{i,j+2}$ 等。
6. 每个 GPU 重复此过程，直到所有 $K$、$V$ 分片都已接收，并且所有部分注意力矩阵 $A_{i,*}$ 都已计算。

Accelerate 通过 accelerator.maybe_context_parallel 装饰器实现此功能，该装饰器也在 Accelerate 示例脚本中展示。您还可以在我们的 CP 概念指南中了解其工作原理和限制。

与 TP 类似，在 Accelerate 中，CP 大小通过 `ParallelismConfig` 中的 `cp_size` 配置，而在 Axolotl 中，您可以使用 `context_parallel_size` 配置字段。

ND 并行

在多节点设置中，FSDP 等数据并行技术将整个网络拓扑视为沿单个维度存在。您可能会发现这种方法在多种原因下受到限制：

• 当扩展到更多节点时，FSDP 的集体操作会受到节点间延迟的瓶颈，导致训练速度过慢。
• 如前所述，大型模型的解码器层可能无法适应 GPU 内存，或者即使处于分片状态，也可能太大而无法执行前向传播。
• 可能无法达到理想的批处理大小——批处理可能太大，纯数据并行无法有效处理，或者由于模型大小的内存限制而太小。

为了解决其中一些问题，我们可以将多节点集群视为具有二维拓扑：设备之间沿一个轴进行快速节点内通信，而沿另一个轴进行相对较慢的节点间通信。让我们考虑如何组合我们迄今为止介绍的并行技术来利用这一点。

混合分片数据并行

混合分片数据并行（HSDP）是一种二维并行，它在节点内执行 FSDP，并在节点间执行 DP——也就是说，模型在每个节点之间复制，并在每个节点内使用 FSDP 进行分片。这使得 FSDP 较高的通信开销可以利用更快的节点内链路，而 DP 将较慢的节点间通信开销最小化到单个梯度同步步骤。如果您遇到问题 1，并希望以增加内存使用为代价来加速训练，您可能会考虑这种方法。

重要的是要注意，我们可以自由配置我们的 2D 网络拓扑的形状，因为我们不受限于维度与物理节点边界对齐——您可能会在 2 个节点之间应用 FSDP，同时在 2 个节点的组之间复制，这将导致较低的内存使用但较慢的吞吐量，但仍将节点内 FSDP 通信开销减少一半。这是一个我们鼓励您根据您的特定硬件设置和微调需求进行调整的参数。

您可以通过在 Accelerate 的 `ParallelismConfig` 或 Axolotl 的配置字段中同时定义 `dp_shard_size` 和 `dp_replicate_size` 来启用 HSDP。

完全分片数据并行 + 张量并行

正如我们之前提到的，TP 应该在节点内部应用以利用高带宽的节点内通信。因此，将 TP 和 FSDP 结合起来涉及使用 FSDP 在节点间对模型进行分片，并在节点内部使用 TP。在一定程度上，这可能为上述所有三个问题提供一个简洁的解决方案：FSDP 的延迟成本可以减少 8 倍，太大无法在单个设备上容纳的层现在均匀分布在设备上，并且由于每个 TP 组接收相同批次的数据，我们还可以将全局批次大小减少 8 倍。然而，如果这仍然不足，我们将无法增加跨节点的 TP 大小，并且必须考虑替代方法。

在 Accelerate 中，您可以通过在 `ParallelismConfig` 中同时定义 `dp_shard_size` 和 `tp_size` 来结合 TP 和 FSDP，而在 Axolotl 中，您可以添加 `dp_shard_size` 和 `tensor_parallel_size` 这两个配置字段。

完全分片数据并行 + 上下文并行

这是一种结合 FSDP 和 CP 的二维并行策略，虽然它并不常用，因为 CP 已经与 FSDP 结合（关于原因请参见 accelerate 概念指南），但在某些情况下它可能很有用，例如需要大序列长度，因此需要大 `cp_size`。如果这仍然不符合您的内存预算，您可以在此之上应用 FSDP，进一步减少内存使用。

在 Accelerate 中，您可以通过在 `ParallelismConfig` 中同时定义 `dp_shard_size` 和 `cp_size` 来结合 CP 和 FSDP，而在 Axolotl 中，您可以添加 `dp_shard_size` 和 `context_parallel_size` 这两个配置字段。

混合分片数据并行 + 张量并行

在足够大的世界大小下（注意：3D 并行的最小世界大小为 8，但在更大规模下最有效），我们可以考虑将 HSDP 与 TP 结合起来，这样就创建了一个层级结构：DP 首先在节点组之间复制模型，然后 FSDP 在每个组内分片模型，最后 TP 在每个节点内拆分单个层。当您面临上述所有扩展限制时，您可能会考虑这种方法，因为它通过在内存使用和吞吐量之间进行权衡，提供了最大的灵活性来适应您的特定训练设置。

在 Accelerate 中，您可以通过在 `ParallelismConfig` 中同时定义 `dp_shard_size`、`dp_replicate_size` 和 `tp_size` 来结合 HSDP 和 TP。类似地，在 Axolotl 中，您可以添加 `dp_shard_size`、`dp_replicate_size` 和 `tensor_parallel_size` 这三个配置字段。

使用注意事项

我们没有涵盖其他并行组合方式，例如使用 HSDP + TP + CP 的 4D 并行，但它们与我们已经涵盖的技术操作方式非常相似。最重要的是，我们鼓励您尝试不同的技术和配置——这是您掌握不同内存/吞吐量权衡方式的最佳途径。

以下是一些您在分布式设置中可能会觉得有用的额外提示：

• 当使用 FSDP 并处理单个设备无法容纳的过大模型时，启用 CPU RAM 高效加载和分片状态字典检查点技术至关重要。您可以通过 Accelerate 的 FullyShardedDataParallelPlugin 中的 cpu_ram_efficient_loading 和 state_dict_type 参数来启用此功能，

  fsdp2_plugin = FullyShardedDataParallelPlugin(
      fsdp_version=2,
      auto_wrap_policy="transformer_based_wrap",
      transformer_cls_names_to_wrap=["LlamaDecoderLayer"],
      state_dict_type="SHARDED_STATE_DICT", 
      cpu_ram_efficient_loading=True
  )
  

  或者通过 Axolotl 中 fsdp_config 里的 cpu_ram_efficient_loading 和 state_dict_type 配置字段。

  fsdp_version: 2
  fsdp_config:
    auto_wrap_policy: TRANSFORMER_BASED_WRAP
    transformer_layer_cls_to_wrap: LlamaDecoderLayer
    state_dict_type: SHARDED_STATE_DICT
    cpu_ram_efficient_loading: True
  

• 训练期间使用的总批次大小对训练稳定性、内存使用和数据吞吐量起着重要作用。当使用 DP 和/或 FSDP 时，有效批次大小计算如下：

  effective_batch_size = micro_batch_size * gradient_accumulation_steps * dp_world_size.

  其中 dp_world_size = (dp_shard_size * dp_replicate_size) / tp_size。您可以通过增加训练循环中的总微批次大小或梯度累积步数，或在 Axolotl 中设置 micro_batch_size 和 gradient_accumulation_steps 配置字段，或通过增加更多 GPU 来增加总 dp_world_size 来增大批次大小。如前所述，这将施加一个 dp_world_size 的*最小*总批次大小——当使用纯 DP/FSDP 时，这将是您的总世界大小，如果这过高，减少总批次大小的唯一方法是引入张量并行。最后，在 GPU 数量固定且内存受限的情况下，我们建议增加 gradient_accumulation_steps 而不是 micro_batch_size 以实现更大的有效批次大小，反之亦然。

• 相应地，当您的有效批次大小因引入数据并行而增加时，您应该缩放学习率以保持训练稳定性。常见的方法包括线性缩放 scaled_lr = base_lr * (effective_batch_size / base_batch_size) 或平方根缩放 scaled_lr = base_lr * sqrt(effective_batch_size / base_batch_size)。

• 即使使用并行策略，如果内存限制仍然存在，梯度检查点可以通过计算换内存的方式提供额外的内存节省。在前向传播期间，只有一部分激活被保存在内存中（通常在 Transformer 块边界），并且中间激活在反向传播期间重新计算。此技术与上述所有并行策略无缝协作。在 Accelerate 中，您可以通过在 FullyShardedDataParallelPlugin 中设置 activation_checkpointing=true 来启用它。

  fsdp2_plugin = FullyShardedDataParallelPlugin(
      fsdp_version=2,
      auto_wrap_policy="transformer_based_wrap",
      transformer_cls_names_to_wrap=["LlamaDecoderLayer"],
      state_dict_type="SHARDED_STATE_DICT", 
      cpu_ram_efficient_loading=True,
      activation_checkpointing=True
  )
  

  在 Axolotl 中也类似。

  fsdp_version: 2
  fsdp_config:
    auto_wrap_policy: TRANSFORMER_BASED_WRAP
    transformer_layer_cls_to_wrap: LlamaDecoderLayer
    state_dict_type: SHARDED_STATE_DICT
    cpu_ram_efficient_loading: True
    activation_checkpointing: True
  

  请注意，梯度检查点通常会因激活重新计算而使训练时间增加约 20-30%，但可以将激活内存减少 60-80%，这使得它在训练非常大的模型或使用长序列长度时特别有价值。