在 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")

编辑并运行代码