使用 Accelerator 的 8-bit Adam

你想用 8-bit Adam 客製化訓練迴圈，以降低模型的記憶體需求。請使用 8-bit Adam 準備訓練迴圈。

假設已經定義好 8-bit Adam 最佳化器 adam_bnb_optim。其他訓練物件也已定義：model、train_dataloader、lr_scheduler 和 accelerator。

本練習屬於課程

使用 PyTorch 高效訓練 AI 模型

檢視課程

練習說明

準備 8-bit Adam 最佳化器以進行分散式訓練。
使用最佳化器更新模型參數。
使用最佳化器將梯度歸零。

動手互動練習

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

# Prepare the 8-bit Adam optimizer for distributed training
model, adam_bnb_optim, train_dataloader, lr_scheduler = accelerator.prepare(model, ____, train_dataloader, lr_scheduler)

for batch in train_dataloader:
    inputs, targets = batch["input_ids"], batch["labels"]
    outputs = model(inputs, labels=targets)
    loss = outputs.loss
    accelerator.backward(loss)
    # Update the model parameters
    adam_bnb_optim.____()
    lr_scheduler.step()
    # Zero the gradients
    adam_bnb_optim.____()
    print(f"Loss = {loss}")

編輯並執行程式碼