huggingfaceのaccelerateを使って訓練時のCUDA out of memoryを回避する

# はじめに

学習スクリプトを実行しているときにGPUにメモリが乗り切らなくて`CUDA out of memory`で処理が落ちてしまい、学習スクリプトを最初から実行し直すハメになることがよくあります。

特に自然言語処理とかだと、batch毎に最大系列長に合わせて短い系列をpaddingするような処理をしている場合、毎ステップで必要なGPUメモリが変化するため、バッチサイズを大きく設定していると1エポック終わるまで不安で仕方ありません。

さらにTransformerベースのアーキテクチャーを使っている場合は、消費メモリが系列長に対して2乗のオーダーなので、ちょっと長い系列長のデータがあったら想定以上にメモリを消費して溢れてしまうケースとかもよくあるんじゃないでしょうか。

huggingfaceのaccelerateというライブラリ内の`find_executable_batch_size`という機能を使えば、そんな課題感を解決できるようなので、accelerateの簡単な紹介もしつつ、本当に`CUDA out of memory`を回避できるのか実装して確かめてみました。

# accelerateとは

- 公式のリファレンスは[こちら](https://huggingface.co/docs/accelerate/index)

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ファイルはこんな感じでした。

```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
```

### 通常の学習コード

まずはよくある通常の学習コード（と自分は思っている）を書きます。毎エポック検証データによるメトリクスも算出するようにしています。


```python: 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クラスはそれぞれ以下のように書いてます。

```python: ./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
```

```python: ./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](https://github.com/huggingface/accelerate/issues/24)にあるように、accelerateで分散学習するときは、`Accelerator`のインスタンス作成時に`find_unused_parameters=True`を設定した`DistributedDataParallelKwargs`を渡す必要があるようです。

その点も踏まえながら、最初の`train.py`をaccelerateが使えるように書き換えるとこんな感じです。
（便宜上、accelerate用に追加した部分の行頭に`+`を書いてます。下記のコードを実行されるときは行頭の`+`は削除してから実行してください。）

```python: 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を回避する

ここからが本題です。

[こちら](https://huggingface.co/docs/accelerate/usage_guides/memory)の公式ドキュメントに記載されているようにdataloaderを宣言するところらへんから学習ループの処理までまるごと内部関数の中に入れてしまいます。

`find_executable_batch_size`の挙動としては、最初に指定したbatch_sizeでcuda out of memoryが発生した場合はbatch_sizeを半分に減らして、再度`find_executable_batch_size`でデコレートされた関数が実行される、それを学習が終了するまで繰り返す、という流れのようです。

上記のaccelerateで動作するコードに`find_executable_batch_size`を追加したコードは以下のような感じです。公式ドキュメント通りに実装しているだけです。

```python: 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が表示され続けます。

```python: 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,
)
```

# おわりに

とても便利だと思いました。