Accelerate 文档

梯度同步

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

梯度同步

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 模型的 forwardbackward 方法中。但这如何会导致你的代码速度变慢呢?

在 DDP(分布式数据并行)中,进程执行和运行的特定顺序在特定点是预期的,并且这些顺序也必须大致同时发生,然后才能继续进行。

最直接的例子是当你通过 optimizer.step() 更新模型参数时。在没有梯度累积的情况下,模型的所有实例都需要在继续处理下一批数据之前,计算、整理和更新它们的梯度。当执行梯度累积时,你累积 n 个损失梯度,并在达到 n 批次之前跳过 optimizer.step()。由于所有训练进程只需要在调用 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.accumulateaccelerator.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.01秒 2.13±0.08秒 0.91±0.11秒 0.91±0.11秒
单节点 0.50±0.01秒 0.50±0.01秒 0.41±0.015秒 0.41±0.015秒

正如你所看到的,如果你不注意如何设置梯度同步,你可能会在训练期间获得超过 2 倍的减速!

如果你担心确保一切都正确完成,我们强烈建议使用 accumulate() 函数,并将 gradient_accumulation_stepsgradient_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(470 亿参数)。我们看到,即使对于适度的 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 上更新