梯度同步

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` 个损失梯度，并跳过 `optimizer.step()` 直到达到 `n` 个批次。由于所有训练进程只需要在 `optimizer.step()` 被调用时同步，不对你的训练步骤做任何修改，这种不必要的进程间通信可能会导致显著的减速。

你如何避免这种开销？

解决减速问题

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

在此上下文管理器下，当调用 `.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 示例脚本的修改版
每个 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` 意味着由于额外的数据同步，*将会出现减速*，正如本指南前面部分所解释的。

< > 在 GitHub 上更新