Zenn
🤗

HuggingfaceのAccelerateの紹介

2021/09/27に公開
Python
ディープラーニング
機械学習
tech

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