將 Accelerate 新增到您的程式碼中

每個分散式訓練框架都有自己的工作方式，這可能需要編寫大量自定義程式碼以使其適應您的 PyTorch 訓練程式碼和訓練環境。Accelerate 提供了一種友好的方式來與這些分散式訓練框架對接，而無需學習每個框架的具體細節。Accelerate 會為您處理這些細節，讓您可以專注於訓練程式碼，並將其擴充套件到任何分散式訓練環境中。

在本教程中，您將學習如何使用 Accelerate 來適配您現有的 PyTorch 程式碼，並輕鬆地開始在分散式系統上進行訓練！您將從一個基本的 PyTorch 訓練迴圈開始（假設所有訓練物件如 model 和 optimizer 都已經設定好），然後逐步將 Accelerate 整合進去。

device = "cuda"
model.to(device)

for batch in training_dataloader:
    optimizer.zero_grad()
    inputs, targets = batch
    inputs = inputs.to(device)
    targets = targets.to(device)
    outputs = model(inputs)
    loss = loss_function(outputs, targets)
    loss.backward()
    optimizer.step()
    scheduler.step()

Accelerator

Accelerator 是適配您的程式碼以使用 Accelerate 的核心類。它瞭解您正在使用的分散式設定，例如不同程序的數量和您的硬體型別。這個類還提供了許多必要的方法，使您的 PyTorch 程式碼能夠在任何分散式訓練環境中工作，並在裝置間管理和執行程序。

因此，您應該總是在指令碼的開始匯入並建立一個 Accelerator 例項。

from accelerate import Accelerator

accelerator = Accelerator()

Accelerator 也知道將您的 PyTorch 物件移動到哪個裝置上，所以建議讓 Accelerate 來為您處理這件事。

- device = "cuda"
+ device = accelerator.device
  model.to(device)

準備 PyTorch 物件

接下來，您需要為分散式訓練準備您的 PyTorch 物件（模型、最佳化器、排程器等）。prepare() 方法負責將您的模型放置在適合您訓練設定的容器中（如單 GPU 或多 GPU），調整最佳化器和排程器以使用 Accelerate 的 AcceleratedOptimizer 和 AcceleratedScheduler，並建立一個可以在程序間分片的新資料載入器。

PyTorch 物件的返回順序與傳入時的順序相同。

model, optimizer, training_dataloader, scheduler = accelerator.prepare(
    model, optimizer, training_dataloader, scheduler
)

訓練迴圈

最後，移除訓練迴圈中對輸入和目標的 to(device) 呼叫，因為 Accelerate 的 DataLoader 類會自動將它們放置在正確的裝置上。您還應該用 Accelerate 的 backward() 方法替換常規的 backward() 傳遞，它會為您縮放梯度並根據您的分散式設定（例如，DeepSpeed 或 Megatron）使用相應的 backward() 方法。

-   inputs = inputs.to(device)
-   targets = targets.to(device)
    outputs = model(inputs)
    loss = loss_function(outputs, targets)
-   loss.backward()
+   accelerator.backward(loss)

將所有內容整合在一起，您的新 Accelerate 訓練迴圈現在應該看起來像這樣！

from accelerate import Accelerator
accelerator = Accelerator()

device = accelerator.device
model, optimizer, training_dataloader, scheduler = accelerator.prepare(
    model, optimizer, training_dataloader, scheduler
)

for batch in training_dataloader:
    optimizer.zero_grad()
    inputs, targets = batch
    outputs = model(inputs)
    loss = loss_function(outputs, targets)
    accelerator.backward(loss)
    optimizer.step()
    scheduler.step()

訓練特性

Accelerate 提供了額外的功能——比如梯度累積、梯度裁剪、混合精度訓練等等——您可以將它們新增到您的指令碼中以改善訓練過程。讓我們來探討這三個特性。

梯度累積

梯度累積允許您透過在更新權重之前累積多個批次的梯度來使用更大的批次大小進行訓練。這對於克服記憶體限制很有用。要在 Accelerate 中啟用此功能，請在 Accelerator 類中指定 gradient_accumulation_steps 引數，並向您的指令碼中新增 accumulate() 上下文管理器。

+ accelerator = Accelerator(gradient_accumulation_steps=2)
  model, optimizer, training_dataloader = accelerator.prepare(model, optimizer, training_dataloader)

  for input, label in training_dataloader:
+     with accelerator.accumulate(model):
          predictions = model(input)
          loss = loss_function(predictions, label)
          accelerator.backward(loss)
          optimizer.step()
          scheduler.step()
          optimizer.zero_grad()

梯度裁剪

梯度裁剪是一種防止“梯度爆炸”的技術，Accelerate 提供了

• clipgrad_value() 用於將梯度裁剪到最小值和最大值之間
• clipgrad_norm() 用於將梯度歸一化到特定值

混合精度

混合精度透過使用像 fp16（半精度）這樣的低精度資料型別來計算梯度，從而加速訓練。為了獲得 Accelerate 的最佳效能，損失計算應在模型內部進行（就像在 Transformers 模型中一樣），因為在模型外部的計算會以全精度進行。

在 Accelerator 中設定要使用的混合精度型別，然後使用 autocast() 上下文管理器來自動將值轉換為指定的資料型別。

+ accelerator = Accelerator(mixed_precision="fp16")
+ with accelerator.autocast():
      loss = complex_loss_function(outputs, target)

儲存和載入

Accelerate 也可以在訓練完成後儲存和載入模型，或者您也可以儲存模型和最佳化器的狀態，這對於恢復訓練很有用。

模型

一旦所有程序完成，請在儲存模型之前使用 unwrap_model() 方法解包模型，因為 prepare() 方法已將您的模型包裝成了適合分散式訓練的介面。如果您不解包模型，儲存模型的狀態字典也會儲存來自更大模型的任何潛在額外層，您將無法將權重載入回您的基礎模型。

您應該使用 save_model() 方法來解包並儲存模型的狀態字典。該方法還可以將模型儲存為分片檢查點或 safetensors 格式。

accelerator.wait_for_everyone()
accelerator.save_model(model, save_directory)

from transformers import AutoModel

unwrapped_model = accelerator.unwrap_model(model)
unwrapped_model.save_pretrained(
    "path/to/my_model_directory",
    is_main_process=accelerator.is_main_process,
    save_function=accelerator.save,
)

model = AutoModel.from_pretrained("path/to/my_model_directory")

unwrapped_model = accelerator.unwrap_model(model)
path_to_checkpoint = os.path.join(save_directory,"pytorch_model.bin")
unwrapped_model.load_state_dict(torch.load(path_to_checkpoint))

狀態

在訓練過程中，您可能希望儲存模型、最佳化器、隨機數生成器以及可能的學習率排程器的當前狀態，以便在同一指令碼中恢復它們。您應該向指令碼中新增 save_state() 和 load_state() 方法來儲存和載入狀態。

要進一步自定義透過 save_state() 儲存狀態的位置和方式，請使用 ProjectConfiguration 類。例如，如果啟用了 automatic_checkpoint_naming，則每個儲存的檢查點都會儲存在 Accelerator.project_dir/checkpoints/checkpoint_{checkpoint_number}。

任何其他需要儲存的有狀態項都應使用 register_for_checkpointing() 方法進行註冊，以便它們可以被儲存和載入。傳遞給此方法以儲存的每個物件都必須具有 load_state_dict 和 state_dict 函式。