梯度同步

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` 個損失梯度，並跳過 `optimizer.step()` 直到達到 `n` 個批次。由於所有訓練程序只需要在 `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.01s	2.13±0.08s	0.91±0.11s	0.91±0.11s
單節點	0.50±0.01s	0.50±0.01s	0.41±0.015s	0.41±0.015s

如你所見，如果你不注意梯度同步的設定，訓練期間可能會出現超過 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（47B 引數）。我們看到，即使對於適度的 `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 上更新