Accelerate 文档
在 Accelerate 中使用本地 SGD
并获得增强的文档体验
开始使用
在 Accelerate 中使用本地 SGD
本地 SGD 是一种分布式训练技术,其中梯度不会在每个步骤都进行同步。因此,每个进程都会更新自己版本的模型权重,在经过给定数量的步骤后,这些权重通过所有进程求平均值进行同步。这提高了通信效率,并且可以显著加快训练速度,尤其是在计算机缺少像 NVLink 这样的更快互连时。与梯度累积(其中提高通信效率需要增加有效批次大小)不同,本地 SGD 不需要更改批次大小或学习率/调度。但是,如果需要,本地 SGD 也可以与梯度累积结合使用。
在本教程中,您将看到如何快速设置本地 SGD Accelerate。与标准的 Accelerate 设置相比,这只需要两行额外的代码。
这个例子将使用一个非常简化的 PyTorch 训练循环,它每两个批次执行一次梯度累积
device = "cuda"
model.to(device)
gradient_accumulation_steps = 2
for index, batch in enumerate(training_dataloader):
inputs, targets = batch
inputs = inputs.to(device)
targets = targets.to(device)
outputs = model(inputs)
loss = loss_function(outputs, targets)
loss = loss / gradient_accumulation_steps
loss.backward()
if (index + 1) % gradient_accumulation_steps == 0:
optimizer.step()
scheduler.step()
optimizer.zero_grad()
将其转换为 Accelerate
首先,将前面显示的代码转换为使用 Accelerate,既不使用 LocalSGD 也不使用梯度累积帮助器。
+ from accelerate import Accelerator
+ accelerator = Accelerator()
+ model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+ model, optimizer, training_dataloader, scheduler
+ )
for index, batch in enumerate(training_dataloader):
inputs, targets = batch
- inputs = inputs.to(device)
- targets = targets.to(device)
outputs = model(inputs)
loss = loss_function(outputs, targets)
loss = loss / gradient_accumulation_steps
+ accelerator.backward(loss)
if (index+1) % gradient_accumulation_steps == 0:
optimizer.step()
scheduler.step()
让 Accelerate 处理模型同步
现在剩下的就是让 Accelerate 为我们处理模型参数同步和梯度累积。为简单起见,我们假设我们需要每 8 个步骤同步一次。这通过添加一条 with LocalSGD
语句和在每个优化器步骤后调用一次 local_sgd.step()
来实现。
+local_sgd_steps=8
+with LocalSGD(accelerator=accelerator, model=model, local_sgd_steps=8, enabled=True) as local_sgd:
for batch in training_dataloader:
with accelerator.accumulate(model):
inputs, targets = batch
outputs = model(inputs)
loss = loss_function(outputs, targets)
accelerator.backward(loss)
optimizer.step()
scheduler.step()
optimizer.zero_grad()
+ local_sgd.step()
在底层,本地 SGD 代码禁用了自动梯度同步(但累积仍然按预期工作!)。取而代之的是,它每 local_sgd_steps
步(以及训练循环结束时)对模型参数进行平均。
局限性
当前的实现仅适用于基本的多 GPU(或多 CPU)训练,而不支持例如 DeepSpeed。
参考文献
尽管我们不知道这种简单方法的真正起源,但本地 SGD 的思想已经很古老,至少可以追溯到
Zhang, J., De Sa, C., Mitliagkas, I., & Ré, C. (2016). Parallel SGD: When does averaging help?. arXiv preprint arXiv:1606.07365.
我们将本地 SGD 一词归功于以下论文(但我们可能没有意识到更早的参考文献)。
Stich, Sebastian Urban. “Local SGD Converges Fast and Communicates Little.” ICLR 2019-International Conference on Learning Representations. No. CONF. 2019.
< > 在 GitHub 上更新