使用 Accelerator 的 Adafactor

你已經用 Trainer 搭配 Adafactor 做出概念驗證，成功以較低記憶體需求訓練語言翻譯模型。現在你想改用 Accelerator 來自訂訓練迴圈。請建立使用 Adafactor 的訓練迴圈！

compute_optimizer_size() 函式已預先定義。部分訓練物件已預先載入：model、train_dataloader 與 accelerator。Adafactor 已從 torch.optim 預先匯入。

本練習屬於課程

使用 PyTorch 高效訓練 AI 模型

檢視課程

練習說明

定義 optimizer 時，將模型參數傳入 Adafactor。
傳入 optimizer 的 state 並印出大小。

動手互動練習

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

# Pass the model parameters to Adafactor
optimizer = Adafactor(params=____.____())
model, optimizer, train_dataloader = accelerator.prepare(model, optimizer, train_dataloader)

for batch in train_dataloader:
    inputs, targets = batch["input_ids"], batch["labels"]
    outputs = model(inputs, labels=targets)
    loss = outputs.loss
    accelerator.backward(loss)
    optimizer.step()
    optimizer.zero_grad()

# Pass in the optimizer state
total_size_megabytes, total_num_elements = compute_optimizer_size(optimizer.state.____())
print(f"Number of optimizer parameters: {total_num_elements:,}\nOptimizer size: {total_size_megabytes:.0f} MB")

編輯並執行程式碼