はじめに
学習スクリプトを実行しているときにGPUにメモリが乗り切らなくてCUDA out of memoryで処理が落ちてしまい、学習スクリプトを最初から実行し直すハメになることがよくあります。
特に自然言語処理とかだと、batch毎に最大系列長に合わせて短い系列をpaddingするような処理をしている場合、毎ステップで必要なGPUメモリが変化するため、バッチサイズを大きく設定していると1エポック終わるまで不安で仕方ありません。
さらにTransformerベースのアーキテクチャーを使っている場合は、消費メモリが系列長に対して2乗のオーダーなので、ちょっと長い系列長のデータがあったら想定以上にメモリを消費して溢れてしまうケースとかもよくあるんじゃないでしょうか。
huggingfaceのaccelerateというライブラリ内のfind_executable_batch_sizeという機能を使えば、そんな課題感を解決できるようなので、accelerateの簡単な紹介もしつつ、本当にCUDA out of memoryを回避できるのか実装して確かめてみました。
accelerateとは
- 公式のリファレンスはこちら
accelerateはTPU、GPU、CPUでの実行を同じコードで記述できるライブラリで、他にもgradient accumulation stepsやfp16などを使った学習を簡単に記述できる便利なライブラリです。
よくある既存の学習コードをちょこっと変更するだけで簡単に使えるようなので、まずはaccelerateを使ってmulti GPUの学習コードを書いてみようと思います。
accelerateの設定
accelerateはpipで簡単にインストールできます。
pip install accelerate
現在のバージョンはaccelerate==0.16.0でした。本記事もこのバージョンで動作確認しているので、バージョンが更新されたら、本記事のコードが動作しないかもしれませんが、ご了承ください。
pipでインストール後、まずはターミナル上でaccelerate configを実行して、対話形式でaccelerateの設定ファイルを作ります。（設定しなくてもacclerete実行時に引数で渡すこともできますが、設定しておいたほうが実行が楽で便利かと思います。）
対話形式の質問内容と回答例は以下のような感じでした。
In which compute environment are you running?
Please select a choice using the arrow or number keys, and selecting with enter
➔ This machine
AWS (Amazon SageMaker)
------------------------------------------------------
Which type of machine are you using?
Please select a choice using the arrow or number keys, and selecting with enter
No distributed training
multi-CPU
➔ multi-GPU
TPU
MPS
------------------------------------------------------
How many different machines will you use (use more than 1 for multi-node training)? [1]: 1
Do you wish to optimize your script with torch dynamo?[yes/NO]:NO
Do you want to use DeepSpeed? [yes/NO]: NO
Do you want to use FullyShardedDataParallel? [yes/NO]: NO
Do you want to use Megatron-LM ? [yes/NO]: NO
How many GPU(s) should be used for distributed training? [1]:2
What GPU(s) (by id) should be used for training on this machine as a comma-seperated list? [all]:0,1
------------------------------------------------------
Do you wish to use FP16 or BF16 (mixed precision)?
Please select a choice using the arrow or number keys, and selecting with enter
no
➔ fp16
bf16
accelerate configuration saved at /home/username/.cache/huggingface/accelerate/default_config.yaml
上の設定で、手元の自分のマシン1台でGPU2枚（0番と1番）を並列に使います、さらにfp16で学習する設定をしています。
出来上がったyamlファイルはこんな感じでした。
compute_environment: LOCAL_MACHINE
deepspeed_config: {}
distributed_type: MULTI_GPU
downcast_bf16: 'no'
dynamo_backend: 'NO'
fsdp_config: {}
gpu_ids: 0,1
machine_rank: 0
main_training_function: main
megatron_lm_config: {}
mixed_precision: fp16
num_machines: 1
num_processes: 2
rdzv_backend: static
same_network: true
use_cpu: false
現在のaccelerateの設定はaccelerate envで簡単に確認できます。
$ accelerate env
Copy-and-paste the text below in your GitHub issue
- `Accelerate` version: 0.16.0
- Platform: Linux-4.15.0-143-generic-x86_64-with-glibc2.27
- Python version: 3.9.4
- Numpy version: 1.24.1
- PyTorch version (GPU?): 1.13.1+cu117 (True)
- `Accelerate` default config:
- compute_environment: LOCAL_MACHINE
- distributed_type: MULTI_GPU
- mixed_precision: fp16
- use_cpu: False
- dynamo_backend: NO
- num_processes: 2
- machine_rank: 0
- num_machines: 1
- gpu_ids: 0,1
- rdzv_backend: static
- same_network: True
- main_training_function: main
- deepspeed_config: {}
- fsdp_config: {}
- megatron_lm_config: {}
- downcast_bf16: no
設定を変えたいときはもう一回accelerate configを実行すれば良いです。
学習スクリプトをaccelerateで書き換える
今回は自然言語処理における、英語テキストから感情予測をするタスクを扱おうと思います。データはhuggingfaceのdatasetsライブラリを使い、モデルはhuggingfaceのtransformersを使っています。必要に応じてそれらを事前にpipでインストールしておきます。バージョンは以下の通りでした。
$ pip list | grep -e datasets -e transformers
datasets 2.9.0
transformers 4.26.0
通常の学習コード
まずはよくある通常の学習コード（と自分は思っている）を書きます。毎エポック検証データによるメトリクスも算出するようにしています。
train.py
import argparse
from tqdm import tqdm
import numpy as np
from sklearn.metrics import f1_score
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import torch.optim as optim
from transformers.modeling_outputs import ModelOutput
from transformers import AutoTokenizer, AutoModel
import datasets
from src.datacollator import EmotionCollator
from src.models import EmotionClassifier
def main(args):
device = 'cuda:0'
# huggingfaceのdatasetsを使って英語の感情分類タスク用のデータをロードする
emotion_dataset = datasets.load_dataset('emotion')
emotion_labels = ["sadness", "joy", "love", "anger", "fear", "surprise"]
# 英語版BERTを使う
tokenizer = AutoTokenizer.from_pretrained(args.model_name)
model = AutoModel.from_pretrained(args.model_name)
collator = EmotionCollator(tokenizer, max_length=args.max_length)
train_dataloader = DataLoader(emotion_dataset['train'], collate_fn=collator, batch_size=args.batch_size, shuffle=True)
valid_dataloader = DataLoader(emotion_dataset['validation'], collate_fn=collator, batch_size=args.batch_size, shuffle=False)
loss_fct = nn.CrossEntropyLoss()
model = EmotionClassifier(model, len(emotion_labels), loss_fct)
model.to(device)
optimizer = optim.Adam(model.parameters(), lr=args.learning_rate)
# 学習＆検証ループ
for num_epoch in range(args.num_epochs):
batch_losses = []
model.train()
for batch in tqdm(train_dataloader):
optimizer.zero_grad()
batch = {k:v.to(device) for k,v in batch.items()}
outputs = model(**batch)
loss = outputs.loss
loss.backward()
optimizer.step()
batch_losses.append(loss.item())
print(f'Epoch: {num_epoch+1}\tloss: {np.array(batch_losses).mean()}')
model.eval()
all_predictions = []
all_labels = []
for batch in tqdm(valid_dataloader):
with torch.no_grad():
outputs = model(**batch)
predictions = outputs.logits.argmax(dim=1)
all_predictions += predictions.cpu().tolist()
all_labels += batch['labels'].cpu().tolist()
score = f1_score(all_labels, all_predictions, average='macro')
print(f'Epoch: {num_epoch+1}\tscore: {score}')
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--model_name", type=str, default="bert-base-cased")
parser.add_argument("--max_length", type=int, default=512)
parser.add_argument("--batch_size", type=int, default=32)
parser.add_argument("--learning_rate", type=float, default=5e-05)
parser.add_argument("--num_epochs", type=int, default=5)
args = parser.parse_args()
main(args)
モデルの定義やcollatorクラスはそれぞれ以下のように書いてます。
./src/datacollator.py
import torch
class EmotionCollator():
def __init__(self, tokenizer, max_length=512):
self.tokenizer = tokenizer
self.max_length = max_length
def __call__(self, examples):
texts = []
labels = []
for example in examples:
texts.append(example['text'])
labels.append(example['label'])
encoding = self.tokenizer(
texts,
padding=True,
truncation=True,
max_length=self.max_length,
return_tensors='pt'
)
encoding['labels'] = torch.tensor(labels)
return encoding
./src/models.py
import torch
import torch.nn as nn
from transformers.modeling_outputs import ModelOutput
class EmotionClassifier(nn.Module):
def __init__(self, base_model, num_classes, loss_fct=None):
super().__init__()
self.encoder = base_model
self.hidden_size = self.encoder.config.hidden_size
self.linear = nn.Linear(self.hidden_size, num_classes)
self.loss_fct = loss_fct
def compute_loss(self, logits, labels):
return self.loss_fct(logits, labels)
def forward(self, input_ids, attention_mask, token_type_ids, labels):
outputs = self.encoder(
input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids
)
outputs = outputs.last_hidden_state[:, 0, :]
outputs = self.linear(outputs)
loss = self.compute_loss(outputs, labels)
return ModelOutput(
logits=outputs,
loss=loss
)
上のコードだとGPUを1枚だけ使う書き方になってます。これをaccelerateを使ってmulti GPUに対応させてみます。
accelerateで書き換え
上記のようなコードを公式ドキュメントのQuick tour通りに変更すると、以下のようなエラーが出てしまいます。
RuntimeError: Expected to have finished reduction in the prior iteration before starting a new one. This error indicates that your module has parameters that were not used in producing loss. You can enable unused parameter detection by passing the keyword argument `find_unused_parameters=True` to `torch.nn.parallel.DistributedDataParallel`, and by
making sure all `forward` function outputs participate in calculating loss.
こちらのissueにあるように、accelerateで分散学習するときは、Acceleratorのインスタンス作成時にfind_unused_parameters=Trueを設定したDistributedDataParallelKwargsを渡す必要があるようです。
その点も踏まえながら、最初のtrain.pyをaccelerateが使えるように書き換えるとこんな感じです。
（便宜上、accelerate用に追加した部分の行頭に+を書いてます。下記のコードを実行されるときは行頭の+は削除してから実行してください。）
train_with_accelerate.py
import argparse
from tqdm import tqdm
import numpy as np
from sklearn.metrics import f1_score
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import torch.optim as optim
from transformers.modeling_outputs import ModelOutput
from transformers import AutoTokenizer, AutoModel
import datasets
from src.datacollator import EmotionCollator
from src.models import EmotionClassifier
# accelerateをインポート
+ from accelerate import Accelerator
+ from accelerate import DistributedDataParallelKwargs
def main(args):
# https://github.com/huggingface/accelerate/issues/24
+ ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
+ accelerator = Accelerator(kwargs_handlers=[ddp_kwargs])
+ device = accelerator.device
emotion_dataset = datasets.load_dataset('emotion')
emotion_labels = ["sadness", "joy", "love", "anger", "fear", "surprise"]
tokenizer = AutoTokenizer.from_pretrained(args.model_name)
model = AutoModel.from_pretrained(args.model_name)
collator = EmotionCollator(tokenizer, max_length=args.max_length)
train_dataloader = DataLoader(emotion_dataset['train'], collate_fn=collator, batch_size=args.batch_size, shuffle=True)
valid_dataloader = DataLoader(emotion_dataset['validation'], collate_fn=collator, batch_size=args.batch_size, shuffle=False)
loss_fct = nn.CrossEntropyLoss()
model = EmotionClassifier(model, len(emotion_labels), loss_fct)
model.to(device)
optimizer = optim.Adam(model.parameters(), lr=args.learning_rate)
+ model, optimizer, train_dataloader, valid_dataloader = accelerator.prepare(
+ model, optimizer, train_dataloader, valid_dataloader
+ )
# 学習＆検証ループ
for num_epoch in range(args.num_epochs):
batch_losses = []
model.train()
for batch in tqdm(train_dataloader):
optimizer.zero_grad()
outputs = model(**batch)
loss = outputs.loss
+ accelerator.backward(loss)
optimizer.step()
batch_losses.append(loss.item())
# accelerator.printを使うことでメインプロセスのみprintできる
+ accelerator.print(f'Epoch: {num_epoch+1}\tloss: {np.array(batch_losses).mean()}')
model.eval()
all_predictions = []
all_labels = []
for batch in tqdm(valid_dataloader):
with torch.no_grad():
outputs = model(**batch)
predictions = outputs.logits.argmax(dim=1)
# GPU毎に計算されているメトリックス情報をマージしている
+ predictions, labels = accelerator.gather_for_metrics((predictions, batch["labels"]))
all_predictions += predictions.cpu().tolist()
all_labels += labels.cpu().tolist()
score = f1_score(all_labels, all_predictions, average='macro')
+ accelerator.print(f'Epoch: {num_epoch+1}\tscore: {score}')
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--model_name", type=str, default="bert-base-cased")
parser.add_argument("--max_length", type=int, default=512)
parser.add_argument("--batch_size", type=int, default=32)
parser.add_argument("--learning_rate", type=float, default=5e-05)
parser.add_argument("--num_epochs", type=int, default=5)
args = parser.parse_args()
main(args)
既存のコードに少し手を加えるだけでaccelerate用に変更できました。
このスクリプトを実行するときは、accelerate launch {スクリプト名}.pyのように実行すればOKです。例えば以下のような感じです。
accelerate launch train_with_accelerate.py --batch_size 32 --num_epochs 3
accelerateの使い方がなんとなくわかりました。
find_executable_batch_sizeを使ってCUDA out of memoryを回避する
ここからが本題です。
こちらの公式ドキュメントに記載されているようにdataloaderを宣言するところらへんから学習ループの処理までまるごと内部関数の中に入れてしまいます。
find_executable_batch_sizeの挙動としては、最初に指定したbatch_sizeでcuda out of memoryが発生した場合はbatch_sizeを半分に減らして、再度find_executable_batch_sizeでデコレートされた関数が実行される、それを学習が終了するまで繰り返す、という流れのようです。
上記のaccelerateで動作するコードにfind_executable_batch_sizeを追加したコードは以下のような感じです。公式ドキュメント通りに実装しているだけです。
train_with_accelerate_find_batch_size.py
import argparse
from tqdm import tqdm
import numpy as np
from sklearn.metrics import f1_score
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import torch.optim as optim
from transformers.modeling_outputs import ModelOutput
from transformers import AutoTokenizer, AutoModel
import datasets
from src.datacollator import EmotionCollator
from src.models import EmotionClassifier
from accelerate import Accelerator
from accelerate import DistributedDataParallelKwargs # https://github.com/huggingface/accelerate/issues/24
+ from accelerate.utils import find_executable_batch_size
def main(args):
ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
accelerator = Accelerator(kwargs_handlers=[ddp_kwargs])
device = accelerator.device
# huggingfaceのdatasetsを使って英語の感情分類タスク用のデータをロードする
emotion_dataset = datasets.load_dataset('emotion')
emotion_labels = ["sadness", "joy", "love", "anger", "fear", "surprise"]
+ @find_executable_batch_size(starting_batch_size=args.batch_size)
+ def inner_training_loop(batch_size):
+ nonlocal accelerator # Ensure they can be used in our context
+ accelerator.free_memory() # Free all lingering references
# 現在のバッチサイズが何なのかわかりにくいので、確認のため、明示的に表示しておく
+ print(f'batch_size: {batch_size}')
# 英語版BERTを使う
tokenizer = AutoTokenizer.from_pretrained(args.model_name)
model = AutoModel.from_pretrained(args.model_name)
collator = EmotionCollator(tokenizer, max_length=args.max_length)
train_dataloader = DataLoader(emotion_dataset['train'], collate_fn=collator, batch_size=batch_size, shuffle=True)
valid_dataloader = DataLoader(emotion_dataset['validation'], collate_fn=collator, batch_size=batch_size, shuffle=False)
loss_fct = nn.CrossEntropyLoss()
model = EmotionClassifier(model, len(emotion_labels), loss_fct)
model.to(device)
optimizer = optim.Adam(model.parameters(), lr=args.learning_rate)
model, optimizer, train_dataloader, valid_dataloader = accelerator.prepare(
model, optimizer, train_dataloader, valid_dataloader
)
# 学習＆検証ループ
for num_epoch in range(args.num_epochs):
batch_losses = []
model.train()
for batch in tqdm(train_dataloader):
optimizer.zero_grad()
outputs = model(**batch)
loss = outputs.loss
accelerator.backward(loss)
optimizer.step()
batch_losses.append(loss.item())
accelerator.print(f'Epoch: {num_epoch+1}\tloss: {np.array(batch_losses).mean()}')
model.eval()
all_predictions = []
all_labels = []
for batch in tqdm(valid_dataloader):
with torch.no_grad():
outputs = model(**batch)
predictions = outputs.logits.argmax(dim=1)
predictions, labels = accelerator.gather_for_metrics((predictions, batch["labels"]))
all_predictions += predictions.cpu().tolist()
all_labels += labels.cpu().tolist()
score = f1_score(all_labels, all_predictions, average='macro')
accelerator.print(f'Epoch: {num_epoch+1}\tscore: {score}')
+ inner_training_loop()
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--model_name", type=str, default="bert-base-cased")
parser.add_argument("--max_length", type=int, default=512)
parser.add_argument("--batch_size", type=int, default=32)
parser.add_argument("--learning_rate", type=float, default=5e-05)
parser.add_argument("--num_epochs", type=int, default=5)
args = parser.parse_args()
main(args)
動作確認
実際にCUDA out of memoryが起きずに適切なbatch_sizeで実行してくれるのか確かめてみます。
ここでの実験ではRTX A6000 (48GB)を1枚使って実験しています。
最初に与えるbatch_sizeは2048とかなり巨大な値にしてみます。もちろんそのままではCUDA out of memoryで落ちてしまうbatch_sizeです。
accelerate launch train_with_accelerate_find_batch_size.py --batch_size 2048
出力は以下のような感じになりました。（warning等一部不要な出力は省いています。）
CUDA out of memoryで落ちることなく学習が完了していることが確認できます。
batch_sizeも半分にだんだんと減っているようです。注目すべきはbatch_sizeが1024のときで、4step分はout of memoryになることなく実行できているようですが、5step目で処理が終わっています。系列長の長いデータを引いてしまったのでしょう。
No config specified, defaulting to: emotion/split
Found cached dataset emotion
100%|█████████████████████████████████████████████████████████| 3/3 [00:00<00:00, 896.35it/s]
batch_size: 2048
0%| | 0/8 [00:02<?, ?it/s]
batch_size: 1024
31%|█████████████████▊ | 5/16 [00:05<00:12, 1.13s/it]
batch_size: 512
100%|████████████████████████████████████████████████████████| 32/32 [00:13<00:00, 2.29it/s]
Epoch: 1 loss: 1.0747559200972319
100%|██████████████████████████████████████████████████████████| 4/4 [00:00<00:00, 6.35it/s]
Epoch: 1 score: 0.8089526602705339
100%|████████████████████████████████████████████████████████| 32/32 [00:13<00:00, 2.33it/s]
Epoch: 2 loss: 0.25367175647988915
100%|██████████████████████████████████████████████████████████| 4/4 [00:00<00:00, 6.20it/s]
Epoch: 2 score: 0.9027841410826865
100%|████████████████████████████████████████████████████████| 32/32 [00:14<00:00, 2.28it/s]
Epoch: 3 loss: 0.1287951988633722
100%|██████████████████████████████████████████████████████████| 4/4 [00:00<00:00, 6.28it/s]
Epoch: 3 score: 0.9064336892133973
100%|████████████████████████████████████████████████████████| 32/32 [00:13<00:00, 2.30it/s]
Epoch: 4 loss: 0.0992106948979199
100%|██████████████████████████████████████████████████████████| 4/4 [00:00<00:00, 6.19it/s]
Epoch: 4 score: 0.9119660799451724
100%|████████████████████████████████████████████████████████| 32/32 [00:14<00:00, 2.25it/s]
Epoch: 5 loss: 0.07830018864478916
100%|██████████████████████████████████████████████████████████| 4/4 [00:00<00:00, 5.48it/s]
Epoch: 5 score: 0.9184996225474085
今回のデータとGPUならbatch_sizeが512で回ることが確認できたので、次回からは初回でbatch_size=512を渡して学習すればよいでしょう。
Trainerクラスを使っていればauto_find_batch_size=Trueを指定するだけでよい
acceleratorがインストールされている環境で学習のコードをhuggingfaceのTrainerクラスで実装している場合はTrainingArgumentsの引数としてauto_find_batch_size=Trueを指定するだけでbatch_sizeによるCUDA out of memoryを回避できます。
挙動は上記のfind_executable_batch_sizeと同様にbatch_sizeを半分にして再実行を自動で行ってくれているようです。
ちょっと注意点としては、Trainerクラスで学習する際、デフォルトでプログレスバーなりメッセージが色々表示されるかと思いますが、そこではずっとbatch_sizeは初期値で与えたbatch_sizeが表示され続けます。
TrainingArgumentsの実行例
training_args = TrainingArguments(
output_dir='./output',
learning_rate=args.learning_rate,
evaluation_strategy="epoch",
logging_strategy="epoch",
save_strategy="epoch",
save_total_limit=1,
dataloader_num_workers=4,
lr_scheduler_type='constant',
per_device_train_batch_size=args.batch_size,
per_device_eval_batch_size=args.batch_size,
num_train_epochs=args.num_epochs,
remove_unused_columns=False,
fp16=True,
+ auto_find_batch_size=True,
)
おわりに
とても便利だと思いました。