Accelerate 文档
梯度同步
并获得增强的文档体验
开始使用
梯度同步
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 无法自动判断你何时需要这样做,但它们确实提供了一个工具来帮助,即添加到你的模型中的 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.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_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(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
意味着会有减速,这是由于额外的数据同步造成的,正如本指南前面的章节所解释的那样。