在 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 上更新