結合 Accelerator 的梯度檢查點

你正在持續最佳化記憶體使用，以便在你的裝置上訓練翻譯模型。梯度累積已經幫助你有效以較大的批次大小進行訓練。請在此基礎上加入梯度檢查點，以進一步降低模型的記憶體占用。

model、train_dataloader 與 accelerator 已經預先定義。

本練習屬於課程

使用 PyTorch 高效訓練 AI 模型

檢視課程

練習說明

在 model 上啟用梯度檢查點。
設定一個 Accelerator 情境管理器，在 model 上啟用梯度累積。

動手互動練習

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

# Enable gradient checkpointing on the model
____.____()

for batch in train_dataloader:
    with accelerator.accumulate(model):
        inputs, targets = batch["input_ids"], batch["labels"]
        # Get the outputs from a forward pass of the model
        ____ = ____(____, labels=targets)
        loss = outputs.loss
        accelerator.backward(loss)
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()
        print(f"Loss = {loss}")

編輯並執行程式碼