Trainer を使った Adafactor

数十億のパラメータを持つ Transformer モデルを、言語翻訳サービス向けに学習させています。計算リソースへの負荷が大きいため、AdamW と比べてメモリ使用量を抑えられる Adafactor オプティマイザを試すことにしました。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")

コードを編集して実行