使用 Accelerator 前後的訓練迴圈

你想把一個 PyTorch 訓練迴圈改為使用 Accelerator，讓你的語言模型能以更簡潔的方式在 MPRC 句子同義改寫資料集上進行翻譯任務。請更新訓練迴圈，讓模型準備好進行分散式訓練。

部分資料已預先載入：

accelerator 是 Accelerator 的實例
train_dataloader、optimizer、model 與 lr_scheduler 已用 Accelerator 定義並完成準備

本練習屬於課程

使用 PyTorch 高效訓練 AI 模型

檢視課程

練習說明

更新 .to(device) 行，改由 Accelerator 處理裝置放置。
修改梯度計算以使用 Accelerator。

動手互動練習

試著完成這個範例程式碼，體驗一下這個練習。

for batch in train_dataloader:
    optimizer.zero_grad()
    inputs, targets = batch["input_ids"], batch["labels"]
    # Update the lines so Accelerator handles device placement
    inputs = inputs.to(device)
    targets = targets.to(device)
    outputs = model(inputs, labels=targets)
    loss = outputs.loss
    # Modify the gradient computation to use Accelerator
    ____.backward(____)
    optimizer.step()
    lr_scheduler.step()

編輯並執行程式碼