使用 Local SGD 与 Accelerate

Local SGD 是一种分布式训练技术，其中梯度不会在每一步都同步。因此，每个进程都会更新其模型权重的版本，在给定数量的步骤之后，这些权重会通过对所有进程进行平均来进行同步。这提高了通信效率，并且可以显著提高训练速度，尤其是在计算机缺少更快的互连（如 NVLink）的情况下。与梯度累积不同（其中提高通信效率需要增加有效批次大小），Local SGD 不需要更改批次大小或学习率/计划。但是，如果需要，Local SGD 也可以与梯度累积结合使用。

在本教程中，您将了解如何快速设置 Local 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()

在幕后，Local SGD 代码 **禁用** 了自动梯度同步（但累积仍按预期工作！）。相反，它每 local_sgd_steps 步（以及在训练循环结束时）对模型参数进行平均。

局限性

当前的实现仅适用于基本的多 GPU（或多 CPU）训练，不包括，例如 DeepSpeed。。

参考文献

虽然我们不知道这种简单方法的真正起源，但 Local SGD 的想法已经存在很长时间，至少可以追溯到

Zhang, J., De Sa, C., Mitliagkas, I., & Ré, C. (2016). Parallel SGD: When does averaging help?. arXiv preprint arXiv:1606.07365.

我们将 Local SGD 这个术语归功于以下论文（但可能还有更早的参考文献我们不知道）。

Stich, Sebastian Urban. “Local SGD Converges Fast and Communicates Little.” ICLR 2019-International Conference on Learning Representations. No. CONF. 2019.

< > 在 GitHub 上更新