Accelerator を使った 8-bit Adam
8-bit Adam を使ってモデルのメモリ使用量を削減するため、カスタムのトレーニングループを作成したいと思います。8-bit Adam を使ったトレーニングループを準備しましょう。
8-bit Adam オプティマイザーは adam_bnb_optim として定義済みです。また、model、train_dataloader、lr_scheduler、accelerator といったトレーニングに必要なオブジェクトもすでに定義されています。
この演習はコースの一部です
PyTorch による効率的な AI モデルトレーニング
演習の手順
- 8-bit Adam オプティマイザーを分散トレーニング用に準備します。
- オプティマイザーを使ってモデルのパラメーターを更新します。
- オプティマイザーを使って勾配をゼロにリセットします。
実践的なインタラクティブ演習
このサンプルコードを完成させて、この演習に挑戦してみましょう。
# Prepare the 8-bit Adam optimizer for distributed training
model, adam_bnb_optim, train_dataloader, lr_scheduler = accelerator.prepare(model, ____, train_dataloader, lr_scheduler)
for batch in train_dataloader:
inputs, targets = batch["input_ids"], batch["labels"]
outputs = model(inputs, labels=targets)
loss = outputs.loss
accelerator.backward(loss)
# Update the model parameters
adam_bnb_optim.____()
lr_scheduler.step()
# Zero the gradients
adam_bnb_optim.____()
print(f"Loss = {loss}")