梯度同步

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 中，当你调用 [~Accelerator.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 上下文管理器。

在这个上下文管理器下，PyTorch 在调用 .backward() 时会跳过梯度同步，而第一次在该上下文管理器外部调用 .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 被替换为 [~Accelerator.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()

正如你所预期的，[~Accelerator.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：仅在 backward 调用周围使用 no_sync，而不是 forward
• 正确使用 no_sync：正确使用 no_sync 模式
• 使用 accumulate：正确使用 [~Accelerator.accumulate]

以下是每种配置在单节点和双节点设置下遍历 29 批数据时的每批平均秒数：

	Baseline	no_sync improperly	no_sync	accumulate
Multi-Node	2±0.01s	2.13±0.08s	0.91±0.11s	0.91±0.11s
Single Node	0.50±0.01s	0.50±0.01s	0.41±0.015s	0.41±0.015s

正如你所见，如果你在设置梯度同步时不够小心，训练过程中可能会出现超过2倍的减速！

如果你担心是否一切设置得当，我们强烈建议使用 [~Accelerator.accumulate] 函数，并将 gradient_accumulation_steps 或 gradient_accumulation_plugin 传递给 [Accelerator] 对象，让 Accelerate 为你处理这些细节。

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

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

因此，在使用 FSDP 且内存紧张的情况下，我们建议在 [~utils.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

WARNING

禁用 no_sync 意味着由于额外的数据同步，将会出现减速，如本指南前面部分所述。