使用 Accelerator 進行梯度累積

你正在訓練一個語言模型，透過改寫複雜句子來簡化翻譯，但你的 GPU 記憶體不足。梯度累積可以讓模型以較小、能放進記憶體的批次來模擬較大的批次進行有效訓練。你希望明確撰寫訓練迴圈以看清其結構，因此使用 Accelerator。注意：本練習實際在 CPU 上執行，但在 GPU 上使用的程式碼相同。

model、train_dataloader、optimizer 和 lr_scheduler 都已預先定義。

本練習屬於課程

使用 PyTorch 高效訓練 AI 模型

檢視課程

練習說明

設定 Accelerator() 使用 2 個步驟的梯度累積。
設定一個 Accelerator 情境管理器，替 model 啟用梯度累積。

動手互動練習

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

# Configure Accelerator
accelerator = ____(____=____)
model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(model, optimizer, train_dataloader, lr_scheduler)
for batch in train_dataloader:
    # Set up an Accelerator context manager
    with ____.____(____):
        inputs, targets = batch["input_ids"], batch["labels"]
        outputs = model(inputs, labels=targets)
        loss = outputs.loss
        accelerator.backward(loss)
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()
        print(f"Loss = {loss}")

編輯並執行程式碼