Accelerator を使った勾配累積

複雑な文を言い換えることで翻訳を簡潔にする言語モデルをトレーニングしていますが、GPU のメモリが不足しています。勾配累積（gradient accumulation）を使うと、メモリに収まる小さなバッチを活用して、より大きなバッチで効果的にトレーニングを行えます。トレーニングループの構造を明示的に確認したいため、Accelerator を使用します。なお、この演習は実際には CPU 上で動作しますが、GPU でもコードは同じです。

model、train_dataloader、optimizer、および lr_scheduler はあらかじめ定義されています。

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

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

コースを見る

演習の手順

2 ステップの勾配累積を使用するように Accelerator() を設定します。
model の勾配累積を有効にする Accelerator コンテキストマネージャを設定します。

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

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

# Configure Accelerator
accelerator = ____(____=____)
model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(model, optimizer, train_dataloader, lr_scheduler)
for batch in train_dataloader:
    # Set up an Accelerator context manager
    with ____.____(____):
        inputs, targets = batch["input_ids"], batch["labels"]
        outputs = model(inputs, labels=targets)
        loss = outputs.loss
        accelerator.backward(loss)
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()
        print(f"Loss = {loss}")

コードを編集して実行