Accelerator を使ったローカル SGD

言語翻訳モデルのメモリ使用量を最適化するために、勾配累積（gradient accumulation）と勾配チェックポインティング（gradient checkpointing）を実装しました。しかし、トレーニングがまだ少し遅いため、デバイス間の通信効率を向上させるために、トレーニングループにローカル SGD を追加することにしました。ローカル SGD を使ったトレーニングループを構築しましょう！

model、train_dataloader、accelerator はあらかじめ定義されており、LocalSGD はインポート済みです。

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

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

コースを見る

演習の手順

8 ステップごとに勾配を同期するよう、local_sgd_steps を設定してください。
ローカル SGD コンテキストマネージャーをステップ実行してください。

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

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

# Set up a context manager to synchronize gradients every eight steps
with LocalSGD(accelerator=accelerator, model=model, local_sgd_steps=____, enabled=True) as local_sgd:
    for batch in train_dataloader:
        with accelerator.accumulate(model):
            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()
            # Step the local SGD context manager
            local_sgd.____()

コードを編集して実行