Accelerate を用いた分散学習

モデルが大きくなるにつれて、限られたハードウェアでより大きなモデルを訓練し、訓練速度を大幅に上昇させるための方法として並列処理が浮上してきました。1台のマシンに複数のGPUがあっても、複数のマシンにまたがる複数のGPUがあっても、あらゆるタイプの分散処理セットアップ上でユーザーが簡単に 🤗 Transformers モデルを訓練できるように、 Hugging Face では 🤗 Accelerate ライブラリを作成しました。このチュートリアルでは、PyTorch の訓練ループをカスタマイズして、分散処理環境での訓練を可能にする方法について学びます。

セットアップ

はじめに 🤗 Accelerate をインストールしましょう:

pip install accelerate

そしたらインポートして Accelerator オブジェクトを作成しましょう。Accelerator は分散処理セットアップを自動的に検出し、訓練のために必要な全てのコンポーネントを初期化します。モデルをデバイスに明示的に配置する必要はありません。

>>> from accelerate import Accelerator

>>> accelerator = Accelerator()

Accelerate する準備をしましょう

次に、関連する全ての訓練オブジェクトを prepare メソッドに渡します。これには、訓練と評価それぞれのDataloader、モデル、optimizer が含まれます:

>>> train_dataloader, eval_dataloader, model, optimizer = accelerator.prepare(
...     train_dataloader, eval_dataloader, model, optimizer
... )

Backward

最後に訓練ループ内の loss.backward() を 🤗 Accelerate の backward メソッドで置き換えます：

>>> for epoch in range(num_epochs):
...     for batch in train_dataloader:
...         outputs = model(**batch)
...         loss = outputs.loss
...         accelerator.backward(loss)

...         optimizer.step()
...         lr_scheduler.step()
...         optimizer.zero_grad()
...         progress_bar.update(1)

以下のコードで確認できる通り、訓練ループに4行のコードを追加するだけで分散学習が可能です！

from accelerate import Accelerator
from transformers import AdamW, AutoModelForSequenceClassification, get_scheduler

accelerator = Accelerator()

model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=2)
optimizer = AdamW(model.parameters(), lr=3e-5)

device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
model.to(device)

train_dataloader, eval_dataloader, model, optimizer = accelerator.prepare(
    train_dataloader, eval_dataloader, model, optimizer
)

num_epochs = 3
num_training_steps = num_epochs * len(train_dataloader)
lr_scheduler = get_scheduler(
    "linear",
    optimizer=optimizer,
    num_warmup_steps=0,
    num_training_steps=num_training_steps
)

progress_bar = tqdm(range(num_training_steps))

model.train()
for epoch in range(num_epochs):
    for batch in train_dataloader:
        batch = {k: v.to(device) for k, v in batch.items()}
        outputs = model(**batch)
        loss = outputs.loss
        loss.backward()
        accelerator.backward(loss)

        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()
        progress_bar.update(1)

訓練する

関連するコードを追加したら、スクリプトまたは Colaboratory などのノートブックで訓練を開始します。

スクリプトで訓練する

スクリプトから訓練をしている場合は、設定ファイルを作成・保存するために以下のコマンドを実行してください:

accelerate config

そして次のようにして訓練を開始します:

accelerate launch train.py

ノートブックで訓練する

Colaboratory の TPU の利用をお考えの場合、🤗 Accelerate はノートブック上で実行することもできます。訓練に必要な全てのコードを関数に含め、notebook_launcher に渡してください:

>>> from accelerate import notebook_launcher

>>> notebook_launcher(training_function)

🤗 Accelerate と豊富な機能についてもっと知りたい方はドキュメントを参照してください。