梯度同步

PyTorch 的分布式模块通过系统内所有 GPU 之间的来回通信来工作。这种通信需要时间，在使用 ddp 模块时，确保所有进程了解彼此的状态会在特定的触发点发生。

这些触发点会被添加到 PyTorch 模型中，具体是其 forward() 和 backward() 方法。当模型被 DistributedDataParallel 包装时，就会发生这种情况：

import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel

model = nn.Linear(10, 10)
ddp_model = DistributedDataParallel(model)

在 Accelerate 中，调用 prepare() 并传入模型时，此转换会自动完成。

+ from accelerate import Accelerator
+ accelerator = Accelerator()
  import torch.nn as nn
- from torch.nn.parallel import DistributedDataParallel

  model = nn.Linear(10,10)
+ model = accelerator.prepare(model)

梯度累积中的性能下降

现在你已经了解到，在分布式训练环境中，PyTorch 会为你的 PyTorch 模型的 forward 和 backward 方法添加钩子。但这为何可能导致你的代码运行变慢？

在 DDP（分布式数据并行）中，进程执行和运行的具体顺序在特定节点上是预期固定的，并且这些操作必须在大致相同的时间点完成，之后才能继续执行。

最直接的例子是通过 optimizer.step() 更新模型参数时。如果没有梯度累积，所有模型实例都需要在继续处理下一批数据之前，完成其梯度的计算、收集和更新。在执行梯度累积时，你会累积 n 个损失梯度，并在达到 n 个批次之前跳过 optimizer.step()。由于所有训练进程只需在调用 optimizer.step() 时进行同步，因此在不修改训练步骤的情况下，这种不必要的进程间通信可能会导致显著的性能下降。

如何避免这种开销？

解决减速问题

由于在训练这些批次时跳过了模型参数更新，因此在实际调用 optimizer.step() 之前，无需同步它们的梯度。 PyTorch 无法自动判断何时需要执行此操作，但它提供了一个工具来提供帮助，即转换模型为 DDP 后添加到模型中的 no_sync 上下文管理器。

在此上下文管理器下，当调用 .backward() 时，PyTorch 将跳过梯度同步，而在此上下文管理器之外首次调用 .backward() 时将触发同步。示例如下：

ddp_model, dataloader, optimizer = accelerator.prepare(model, dataloader, optimizer)

for index, batch in enumerate(dataloader):
    inputs, targets = batch
    # Trigger gradient synchronization on the last batch
    if index != (len(dataloader) - 1):
        with ddp_model.no_sync():
            # Gradients only accumulate
            outputs = ddp_model(inputs)
            loss = loss_func(outputs)
            accelerator.backward(loss)
    else:
        # Gradients finally sync
        outputs = ddp_model(inputs)
        loss = loss_func(outputs)
        accelerator.backward(loss)
        optimizer.step()

在 Accelerate 中，为了让该 API 可以在任何训练设备上调用（尽管如果你不在分布式系统中，它可能不会执行任何操作！），ddp_model.no_sync 被替换为 no_sync()，并且工作方式相同：

  ddp_model, dataloader, optimizer = accelerator.prepare(model, dataloader, optimizer)

  for index, batch in enumerate(dataloader):
      inputs, targets = batch
      # Trigger gradient synchronization on the last batch
      if index != (len(dataloader)-1):
-         with ddp_model.no_sync():
+         with accelerator.no_sync(model):
              # Gradients only accumulate
              outputs = ddp_model(inputs)
              loss = loss_func(outputs, targets)
              accelerator.backward(loss)
      else:
          # Gradients finally sync
          outputs = ddp_model(inputs)
          loss = loss_func(outputs)
          accelerator.backward(loss)
          optimizer.step()
          optimizer.zero_grad()

不出所料，accumulate() 函数通过跟踪当前批次编号来封装此条件检查，从而为您提供最终的梯度累积 API：

ddp_model, dataloader, optimizer = accelerator.prepare(model, dataloader, optimizer)

for batch in dataloader:
    with accelerator.accumulate(model):
        optimizer.zero_grad()
        inputs, targets = batch
        outputs = model(inputs)
        loss = loss_function(outputs, targets)
        accelerator.backward(loss)
        optimizer.step()
        optimizer.zero_grad()

因此，在 API 选择上，您应使用 accelerator.accumulate 或 accelerator.no_sync。

到底会有多大的性能下降，以及你容易犯的错误

要设置一个现实的示例，请考虑以下配置：

两个单 GPU T4 节点和一个配备两个 GPU 的节点
每个 GPU 都是 T4，托管在 GCP 上
所使用的脚本是 NLP Example 脚本的修改版本
每个 GPU 的批处理大小为 16，并且每 4 步累积一次梯度

所有脚本均可在此仓库中找到。

如果不注意梯度同步和 GPU 通信，当这些 GPU 在非必要的时间段内相互通信时，可能会浪费大量时间。

具体是多少？

参考：

基线：未使用本文讨论的任何同步方法
no_sync 使用不当：no_sync 应仅包裹 backward 调用，而非 forward
no_sync：正确使用 no_sync 模式
accumulate：正确使用 accumulate()

以下是每种配置在单节点和双节点设置下，对 29 个数据批次进行迭代时每个批次的平均秒数：

	基线	`no_sync` 使用不当	`no_sync`	`accumulate`
多节点	2±0.01s	2.13±0.08s	0.91±0.11s	0.91±0.11s
单节点	0.50±0.01s	0.50±0.01s	0.41±0.015s	0.41±0.015s

如你所见，如果不仔细设置梯度同步，训练过程中可能会出现超过 2 倍的减速！

如果您担心是否一切都已正确完成，我们强烈建议使用 accumulate() 函数，并将 gradient_accumulation_steps 或 gradient_accumulation_plugin 传递给 Accelerator 对象，以便 Accelerate 为您处理此操作。

使用 FSDP 时，no_sync 需要额外的 GPU 内存

请注意，在执行 FSDP 训练时，不同步梯度可能会产生不利影响。正如 torch 中警告的那样，FSDP 的 no_sync 上下文管理器将需要额外的内存。

因此，在使用 FSDP 时，如果遇到内存占用较高的情况，我们建议在 GradientAccumulationPlugin 中将 sync_each_batch 设置为 True，以禁用 no_sync。

参见下方示例，我们在 8 块 A100-80GB GPU 上对 Mixtral（47B 参数）进行微调。可以看到，即使 gradient_accumulation_steps=2 这样较小的值，如果启用了 no_sync，也会很快出现内存溢出（OOM）。同样，这是由于 FSDP 的 no_sync 带来的额外内存开销所致。然而，如果通过 sync_each_batch=True 禁用 no_sync，那么 gradient_accumulation_steps=16 的内存消耗将回落到与 gradient_accumulation_steps=1 相同的水平。

模型	`no_sync` (accum=1)	`no_sync`（accum=2）	`no_sync` 已禁用（accum=16）
mixtral 8x7B	69G	OOM	69G

禁用 no_sync 意味着由于额外的数据同步，会出现性能下降，正如本指南前面的章节所述。

Update on GitHub