使用 Trainer 的 Adafactor

你正在為語言翻譯服務訓練一個擁有數十億參數的 Transformer 模型。它正耗盡你的運算資源，因此你決定嘗試使用 Adafactor 最佳化器，以相較於 AdamW 降低記憶體需求。請為 Adafactor 設定 Trainer！

部分訓練物件已預先載入，包括 model、train_dataset、validation_dataset，以及 compute_metrics。

本練習屬於課程

使用 PyTorch 高效訓練 AI 模型

檢視課程

練習說明

在 TrainingArguments 中指定 Adafactor 作為最佳化器。
傳入最佳化器的狀態以列印其大小。

動手互動練習

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

# Specify Adafactor as an optimizer
training_args = TrainingArguments(output_dir="./results",
                                  evaluation_strategy="epoch",
                                  optim="____")

trainer = Trainer(model=model,
                  args=training_args,
                  train_dataset=train_dataset,
                  eval_dataset=validation_dataset,
                  compute_metrics=compute_metrics)
trainer.train()

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

編輯並執行程式碼