Accelerator を使った Adafactor

Trainer を使った Adafactor の概念実証として、メモリ使用量を抑えた言語翻訳モデルのトレーニングを実証しました。次は、Accelerator を使ってトレーニングループをカスタマイズしてみましょう。Adafactor を使ったトレーニングループを構築してください。

compute_optimizer_size() 関数はあらかじめ定義されています。また、model、train_dataloader、accelerator といったトレーニングオブジェクトも事前に読み込まれています。Adafactor は torch.optim からインポート済みです。

この演習はコースの一部です

PyTorch による効率的な AI モデルトレーニング

コースを見る

演習の手順

optimizer を定義する際に、モデルパラメータを Adafactor に渡してください。
オプティマイザの状態を渡して、サイズを出力してください。

実践的なインタラクティブ演習

このサンプルコードを完成させて、この演習に挑戦してみましょう。

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

コードを編集して実行