隆重推出 🤗 Accelerate
🤗 Accelerate
在任何型別的裝置上執行您的**原始** PyTorch 訓練指令碼。
大多數基於 PyTorch 的高階庫都支援分散式訓練和混合精度，但它們引入的抽象要求使用者在需要自定義底層訓練迴圈時學習新的 API。🤗 Accelerate 是為那些喜歡完全控制訓練迴圈，但不願編寫（和維護）使用分散式訓練（用於單個或多個節點上的多 GPU、TPU 等）或混合精度訓練所需的樣板程式碼的 PyTorch 使用者建立的。未來的計劃包括支援 fairscale、deepspeed、AWS SageMaker 特定的資料並行和模型並行。
它提供了兩件事：一個簡單且一致的 API，抽象了這些樣板程式碼，以及一個啟動器命令，可以輕鬆地在各種設定上執行這些指令碼。
輕鬆整合！
我們先看一個例子
import torch
import torch.nn.functional as F
from datasets import load_dataset
+ from accelerate import Accelerator
+ accelerator = Accelerator()
- device = 'cpu'
+ device = accelerator.device
model = torch.nn.Transformer().to(device)
optim = torch.optim.Adam(model.parameters())
dataset = load_dataset('my_dataset')
data = torch.utils.data.DataLoader(dataset, shuffle=True)
+ model, optim, data = accelerator.prepare(model, optim, data)
model.train()
for epoch in range(10):
for source, targets in data:
source = source.to(device)
targets = targets.to(device)
optimizer.zero_grad()
output = model(source)
loss = F.cross_entropy(output, targets)
- loss.backward()
+ accelerator.backward(loss)
optimizer.step()
只需在任何標準 PyTorch 訓練指令碼中新增五行程式碼，您現在就可以在任何分散式設定下執行該指令碼，無論是否使用混合精度。🤗 Accelerate 甚至可以為您處理裝置放置，因此您可以進一步簡化上述訓練迴圈
import torch
import torch.nn.functional as F
from datasets import load_dataset
+ from accelerate import Accelerator
+ accelerator = Accelerator()
- device = 'cpu'
- model = torch.nn.Transformer().to(device)
+ model = torch.nn.Transformer()
optim = torch.optim.Adam(model.parameters())
dataset = load_dataset('my_dataset')
data = torch.utils.data.DataLoader(dataset, shuffle=True)
+ model, optim, data = accelerator.prepare(model, optim, data)
model.train()
for epoch in range(10):
for source, targets in data:
- source = source.to(device)
- targets = targets.to(device)
optimizer.zero_grad()
output = model(source)
loss = F.cross_entropy(output, targets)
- loss.backward()
+ accelerator.backward(loss)
optimizer.step()
相比之下，為了讓這段程式碼在分散式訓練中執行，需要進行以下更改
+ import os
import torch
import torch.nn.functional as F
from datasets import load_dataset
+ from torch.utils.data import DistributedSampler
+ from torch.nn.parallel import DistributedDataParallel
+ local_rank = int(os.environ.get("LOCAL_RANK", -1))
- device = 'cpu'
+ device = device = torch.device("cuda", local_rank)
model = torch.nn.Transformer().to(device)
+ model = DistributedDataParallel(model)
optim = torch.optim.Adam(model.parameters())
dataset = load_dataset('my_dataset')
+ sampler = DistributedSampler(dataset)
- data = torch.utils.data.DataLoader(dataset, shuffle=True)
+ data = torch.utils.data.DataLoader(dataset, sampler=sampler)
model.train()
for epoch in range(10):
+ sampler.set_epoch(epoch)
for source, targets in data:
source = source.to(device)
targets = targets.to(device)
optimizer.zero_grad()
output = model(source)
loss = F.cross_entropy(output, targets)
loss.backward()
optimizer.step()
這些更改將使您的訓練指令碼適用於多 GPU，但您的指令碼將停止在 CPU 或單 GPU 上工作（除非您到處新增 if 語句）。更令人煩惱的是，如果您想在 TPU 上測試您的指令碼，您將需要更改不同的程式碼行。混合精度訓練也是如此。🤗 Accelerate 的承諾是
- 將您對訓練迴圈的更改保持在最低限度，以便您需要學習的內容儘可能少。
- 讓相同的功能適用於任何分散式設定，因此只需學習一個 API。
它是如何工作的？
要了解該庫在實踐中是如何工作的，讓我們看看我們需要新增到訓練迴圈中的每一行程式碼。
accelerator = Accelerator()
除了提供您將使用的主要物件外，此行還將從環境中分析分散式訓練執行的型別，並執行必要的初始化。您可以透過將 `cpu=True` 或 `fp16=True` 傳遞給此初始化來強制在 CPU 或混合精度訓練。這兩個選項也可以使用指令碼啟動器進行設定。
model, optim, data = accelerator.prepare(model, optim, data)
這是 API 的主要部分，它將準備三種主要型別的物件：模型 (`torch.nn.Module`)、最佳化器 (`torch.optim.Optimizer`) 和資料載入器 (`torch.data.dataloader.DataLoader`)。
模型
模型準備包括將其包裝在適當的容器中（例如 `DistributedDataParallel`）並將其放在適當的裝置上。與常規分散式訓練一樣，您需要解包模型以進行儲存或訪問其特定方法，這可以透過 `accelerator.unwrap_model(model)` 完成。
最佳化器
最佳化器也封裝在一個特殊容器中，該容器將在步驟中執行必要的操作以使混合精度工作。如果其狀態字典非空或從檢查點載入，它還將正確處理狀態字典的裝置放置。
DataLoader
這裡隱藏著大部分魔力。正如您在程式碼示例中看到的那樣，該庫不依賴於 `DistributedSampler`，它實際上適用於您可能傳遞給資料載入器的任何取樣器（如果您曾經不得不編寫自定義取樣器的分散式版本，那麼就不再需要了！）。資料載入器封裝在一個容器中，該容器將只獲取取樣器中與當前程序相關的索引（或者如果使用 `IterableDataset` 則跳過其他程序的批次），並將批次放在適當的裝置上。
為此，Accelerate 提供了一個實用函式，用於在分散式訓練期間同步每個程序上的隨機數生成器。預設情況下，它只同步您的取樣器的 `generator`，因此您的資料增強在每個程序上都會不同，但隨機洗牌將是相同的。如果需要，您當然可以使用此實用程式同步更多的 RNG。
accelerator.backward(loss)
最後一行添加了反向傳播所需的步驟（主要用於混合精度，但其他整合在這裡將需要一些自定義行為）。
評估如何？
評估既可以在所有程序上正常執行，也可以在您只想在主程序上執行時使用方便的測試
if accelerator.is_main_process():
# Evaluation loop
但是您也可以非常輕鬆地使用 Accelerate 執行分散式評估，以下是您需要新增到評估迴圈中的內容
+ eval_dataloader = accelerator.prepare(eval_dataloader)
predictions, labels = [], []
for source, targets in eval_dataloader:
with torch.no_grad():
output = model(source)
- predictions.append(output.cpu().numpy())
- labels.append(targets.cpu().numpy())
+ predictions.append(accelerator.gather(output).cpu().numpy())
+ labels.append(accelerator.gather(targets).cpu().numpy())
predictions = np.concatenate(predictions)
labels = np.concatenate(labels)
+ predictions = predictions[:len(eval_dataloader.dataset)]
+ labels = label[:len(eval_dataloader.dataset)]
metric_compute(predictions, labels)
與訓練一樣，您需要新增一行來準備您的評估資料載入器。然後您只需使用 `accelerator.gather` 在程序之間收集預測和標籤的張量。要新增的最後一行將預測和標籤截斷為資料集中的示例數量，因為準備好的評估資料載入器將返回更多元素以確保每個程序上的批次大小相同。
一個啟動器，統領一切
使用 Accelerate 的指令碼將與您的傳統啟動器完全相容，例如 `torch.distributed.launch`。但記住所有引數有點煩人，當您設定了 4 個 GPU 的例項時，您將使用所有這些 GPU 執行大部分訓練。Accelerate 帶有一個方便的 CLI，分兩步工作
accelerate config
這將觸發一個關於您的設定的小問卷，它將建立一個配置檔案，您可以編輯該檔案以包含所有訓練命令的預設值。然後
accelerate launch path_to_script.py --args_to_the_script
將使用這些預設值啟動您的訓練指令碼。您所要做的就是提供訓練指令碼所需的所有引數。
為了讓這個啟動器更棒，您可以使用它來使用 SageMaker 啟動 AWS 例項。請檢視此指南以瞭解如何操作！
如何參與？
要開始，只需 `pip install accelerate` 或檢視文件以獲取更多安裝選項。
Accelerate 是一個完全開源的專案，您可以在 GitHub 上找到它，檢視其文件或瀏覽我們的基本示例。如果您有任何問題或希望該庫支援的功能，請告訴我們。所有問題，請訪問論壇！
對於更復雜的實際示例，您可以檢視官方的 Transformers 示例。每個資料夾都包含一個利用 Accelerate 庫的 `run_task_no_trainer.py`！
