🤗

HuggingfaceのAccelerateの紹介

2021/09/27に公開

accelerateとは

https://github.com/huggingface/accelerate

いわゆるPyTorchでCPU,GPU,TPUを共通のコードでかけるようにしてくれるもの.
ただし,pytorch-lightningのような高レベルのAPIを持っているわけではなく,どちらかといえば俺俺Trainerをすこし簡単にしてくれるもの.

Run your raw PyTorch training script on any kind of device

まさにこれですね.

まずは,どのように簡単にしてくれるかを見てみましょう

  import torch
  import torch.nn.functional as F
  from datasets import load_dataset
+ from accelerate import Accelerator

- device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+ accelerator = Accelerator()

- model = torch.nn.Transformer().to(device)
+ model = torch.nn.Transformer()
  optimizer = torch.optim.Adam(model.parameters())

  dataset = load_dataset('my_dataset')
  data = torch.utils.data.DataLoader(dataset, shuffle=True)

+ model, optimizer, data = accelerator.prepare(model, optimizer, data)

  model.train()
  for epoch in range(10):
      for source, targets in data:
-         source = source.to(device)
-         targets = targets.to(device)

          optimizer.zero_grad()

          output = model(source)
          loss = F.cross_entropy(output, targets)

-         loss.backward()
+         accelerator.backward(loss)

          optimizer.step()

上のように,基本的に.to(device)はいらず,それらはaccelerator.prepare(*args)でmodel, optimizer, data_loaderをラップすることで対応しています.

上の例ではprepareメソッドにmodel, optimizer, dataを入れていますが,どの順番でもよく何個入れても大丈夫です.
個別にしたい場合には,accelerator.prepare_model等使えます.

AMP

ampにも対応していて,

accelerator = Accelerator(fp16=True)

とするだけでできます.

clip_grad_norm

pytorchにあるtorch.nn.utils.clip_grad_norm_を使いたい場合は以下のようにします.

loss = ...
optimizer.zero_grad()
accelerator.backward(loss)
accelerator.clip_grad_norm_(model.parameters(), max_norm=5, norm_type=2)
optimizer.step()

モデルの保存

DDP等使う場合には,prepareの中でmodelがDistributedDataParallelでラップされているのでモデルを保存は以下のようにunwrapします.

torch.save(accelerator.unwrap_model(model).state_dict(), model_path)
or
accelerator.save(accelerator.unwrap_model(model), model_path)

acceleratorのほうはTPUを使用している場合に,torch_xla.core.xla_model.saveの方で行ってくれます.

マルチノードやマルチGPUを使う

はじめに,以下のコマンドで使用するマシンの状態を答えます.
その後accelerate launch ...でスクリプト実行するとDDPに対応して実行してくれます.

# ノード数やGPU数,FP16を使うかどうかの質問に答える.
$ accelerate config
# その後
$ accelerate launch path_to_script.py --args_for_the_script

このとき,メインプロセスのみで実行したかったり,他プロセスの実行を待ちたい場合(モデルのセーブなど)には以下のメソッドが有用です.

# メインプロセスで実行したい
if accelerator.is_main_process:
    something()
    
# ノードごとに1プロセスの場合
if accelerator.is_local_main_process:
    something()

# 他プロセスの実行を待ちたい
...
accelerator.wait_for_everyone()

また,tqdmなどでは

# tqdmを使う場合(これで1つのバーのみ表示される)
from tqdm.auto import tqdm

progress_bar = tqdm(range(args.max_train_steps), disable=not accelerator.is_local_main_process)

これでバーが複数表示されることがなくなりますね.

Notebook上でのaccelerate launch

colabなどでコマンド実行は煩わしいので,notebook用に関数が用意されています.

from accelerate import notebook_launcher

notebook_launcher(training_function)

感想

ぶっちゃけpytorch-lightningのほうが少ないコードで実装できる気もしますが,それでも俺俺Trainerを書いている皆様,ぜひ使ってみてください.

Discussion