GPUメモリ使用量を削減する


# はじめに

筆者はいわゆる日曜プログラマーで、実行環境は比較的金銭的負担が少なくて済む[Google Colab Pro]を使い、主に画像系を対象にしてきました。一方で、GPT2などNLP系はネットワークが巨大過ぎて[Google Colab Pro]でもGPUメモリ不足にあい、対象にしづらい状況でした。

そんな中、「[How to fine tune VERY large model if it doesn’t fit on your GPU]」という記事に出会い、ここで述べられている技術を使えば[Google Colab Pro]でもかなり大きなネットワークを実行できそうな感触を得ました。また、Hugging Faceの[transformers]ではこれらの技術がコンフィグの設定だけで使用でき、手軽に活用できることがわかりました。そして、実際に10億超のパラメータを持つ[RoBERTaモデル][RoBERTa実装説明]を、GPUメモリ削減技術を使って[Google Colab Pro]で実際に学習することができました。

そこで、この記事では、以下に沿ってGPUメモリ削減技術の有効性を記述してまいります。

- GPUメモリ削減技術の紹介
- GPUメモリ削減技術の効果検証
- 10億超パラメータの[RoBERTaモデル][RoBERTa実装説明]の学習

# GPUメモリ削減技術の紹介

最初に、「[How to fine tune VERY large model if it doesn’t fit on your GPU]」で述べられている技術をご紹介します。

- Gradient Checkpointing
- Gradient Accumulation
- Mixed Precision Training
- Optimizer 8-bit Quantization

なお、個々の技術についてはすでに多くの解説記事も出ていることから、ここでは簡単な紹介にとどめました。[参考資料](#参考資料)にいくつか資料をリストアップしましたので、詳細はそちらをご参照ください。

### Gradient Checkpointing

学習する際、バックプロパゲーションで損失関数の勾配を計算するためにフォワードパスの実行結果を使います。フォワードパス実行結果を覚えておくことで勾配計算を簡単に行うことができますが、その代償としてメモリを消費してしまいます。

Gradient Checkpointingは一部のノードをチェックポイントノードとして指定し、フォワードパスの実行結果をそれらのノードでのみ覚え、他は捨てることでメモリの消費をおさえます。バックプロパゲーションでは、そのノードより前方にあるチェックポイントノードの中で最も近いものからフォワードパス計算を行ってそのノードでのフォワードパス出力値を求め、勾配計算に使います。

これにより、n層からなる深いフィードフォワードニューラルネットワークを学習する場合、メモリ消費量をO(sqrt(n))までおさえられることが論文「[Training Deep Nets with Sublinear Memory Cost][Gradient Checkpointing]」に記載されています。ただし、その代償としてフォワードパスを余分に実行する必要があり、学習時間が長引きます。

### Gradient Accumulation

ミニバッチサイズはハイパーパラメータですので、モデルに応じて適切な値を設定する必要があります。今回対象としている[RoBERTa][RoBERTa実装説明]では大きなミニバッチサイズが良い精度を出すことが論文「[RoBERTa: A Robustly Optimized BERT Pretraining Approach][RoBERTa論文]」で述べられています。

![キャプチャ.JPG](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/629790/61e03c49-e70e-be90-ae73-5deaf785ebbd.jpeg)

一方で、フォワードパス出力値をバックプロパゲーションの際に使うために、ミニバッチサイズに応じたデータをGPUメモリに保持する必要があるので、ミニバッチサイズによってはGPUメモリに乗らない場合があります。

Gradient Accumulationはミニバッチサイズを小さなミニバッチの累積で疑似しようというものです。小さなミニバッチを順次実行する際に損失関数の勾配を累積し、最後の小さなミニバッチ終了時に蓄積した勾配をモデル変数更新のために使用します。

以下の疑似コードは「[Gradient Accumulation in PyTorch]」から引用させていただきました。

![キャプチャ.JPG](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/629790/4742654c-3435-5662-9fb1-20082af2d43b.jpeg)


### Mixed Precision Training

Mixed Precision Trainingは、学習演算処理の中で32ビット浮動小数点演算とそれより小さいフォーマット（FP16, TP16, BF16）での演算処理を使い分けることで、メモリ使用を削減するとともに計算の高速化をはかります。

なお、学習モデルによっては、損失関数の勾配がFP16などでは表現できないレンジに落ち込むためにゼロとなり学習が進まない、というケースがあるようです（下記グラフ参照）。勾配をスケールすることでゼロ化が回避できるようですが、[transformers]では使用するデバイスに応じて`torch_xla.amp.GradScaler()`や`torch.cuda.amp.GradScaler()`などを自動的に呼び出すように実装しているので、使用者はこの問題への対応を意識する必要がありません。

![キャプチャ.JPG](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/629790/81a13717-8949-660a-c892-2b6a39cad191.jpeg)
(「[NVIDIA Train With Mixed Precision]」から抜粋）


### Optimizer 8-bit Quantization


論文「[8-BIT OPTIMIZERS VIA BLOCK-WISE QUANTIZATION]」で提案された８ビットオプティマイザーは、32ビット性能を維持しながらメモリ削減を実現します。

この[論文](https://arxiv.org/pdf/2110.02861.pdf)のTable 1ではさまざまなタスクでの精度とメモリ削減をリスト化しています。例えば、335Mのパラメータをもつ[RoBERTa-Large]では[GLUE]に対して各種オプティマイザーで実行した結果、32ビットに比較して2GBのメモリを削減しながら精度も維持し、かつ高速に実行していることが示されています。

![キャプチャ.JPG](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/629790/d0695bab-1086-38a6-a1b1-9cb2fff1a9c0.jpeg)


# GPUメモリ削減技術の効果検証

上記GPUメモリ削減技術の効果を[RoBERTaモデル][RoBERTa実装説明]で確認したいと思います。

### 検証に使用するJupyter Notebook

「[How to train a new language model from scratch using Transformers and Tokenizers]」で84Mパラメータの[RoBERTaモデル][RoBERTa実装説明]を学習する方法が述べられていますので、これを検証に利用します。[上述の記事][How to train a new language model from scratch using Transformers and Tokenizers]のタイトルの下にある「[Open in Colab]」をクリックすると[Colab][Google Colab Pro]上で動作する[Jupyter Notebook][notebook]を開けますのでコピーしてお使いください。

### 検証環境

筆者が実行したときの環境情報を記載します。なお、ランタイムの仕様で「ハイメモリ」（[Google Colab Pro]以上で指定可能）を使います。

|主なコンポーネント	|コンポーネント情報									|
|--					|--													|				
|OS					|Ubuntu 18.04										|
|Python				|3.7.13												|
|torch				|1.11.0+cu113										|
|transformers		|4.19.0.dev0										|
|tokenizers			|0.12.1												|
|Dataset			|Esperanto portion of the OSCAR corpus from INRIA	|
|CPU				|Intel(R) Xeon(R)CPU @ 2.30GHz						|
|RAM				|25.46 GB											|
|GPU				|Tesla P100-PCIE (16GB GPU Memory)					|

なお、無料のGoogle Colabでは、RAMが12GB程度しか割り当たらないため、使用する[notebook]ではdataset作成でクラッシュしてしまいGPUメモリ削減技術を試すに至りません。

### インストール

Optimizer 8-bit Quantizationを使うために`bitsandbytes-cudaxxx`をインストールします。`xxx`はcudaのバージョンをあてはめます。cudaバージョンは`nvcc -V`コマンドで出力される情報を確認します。例えば、以下ですと111をあてはめることになります。

![キャプチャ.JPG](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/629790/888cf796-e18c-82e5-96e7-fda7f3acb4bb.jpeg)


```
!pip install bitsandbytes-cuda111
```

### GPUメモリ削減技術の設定

[Notebook]でGPUメモリ削減技術を設定するためのパラメータを示します。

|GPUメモリ削減技術				|設定							|
|--								|--								|
|Gradient Checkpointing			|gradient_checkpointing=True,	|
|Gradient Accumulation			|#per_gpu_train_batch_size=64,<BR>per_device_train_batch_size=4,<BR>gradient_accumulation_steps=16,|
|Mixed Precision Training		|fp16=True, # FP16の場合			|
|Optimizer 8-bit Quantization	|optim=OptimizerNames.ADAMW_BNB,|

[Notebook]では`per_gpu_train_batch_size`を使っていますが、deprecatedされ将来削除されるそうなのでコメントアウトし、その代替えとなる`per_device_train_batch_size`を使います。

Gradient Accumulationの設定では、`per_device_train_batch_size`に小さなミニバッチのサイズを指定し、その小さなミニバッチの実行を累積する回数を`gradient_accumulation_steps`で指定します。また、疑似的なミニバッチサイズは`per_device_train_batch_size * gradient_accumulation_steps`で求めることができます。引用している[Notebook]ではミニバッチサイズとして64を指定しているので、疑似的なミニバッチサイズが64となるように上記のとおり設定します。

これらを踏まえ、[Notebook]「Finally, we are all set to initialize our Trainer」のセルを書き換えて実行します。以下は総てのGPUメモリ削減技術を適用する場合のコンフィグです。なお、`output_dir`にはチェックポイントを出力するパスを記述します。

~~~python
import os, glob
from transformers import Trainer, TrainingArguments
from transformers.training_args import OptimizerNames

output_dir="/gpt2/EsperBERTo-all"

training_args = TrainingArguments(
    output_dir=output_dir,
    overwrite_output_dir=True,
    num_train_epochs=5,
    #per_gpu_train_batch_size=64,
    save_steps=10_000,
    save_total_limit=1,
    prediction_loss_only=True,
    fp16=True,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=16,
    gradient_checkpointing=True,
    optim=OptimizerNames.ADAMW_BNB,
)

trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=data_collator,
    train_dataset=dataset,
)

resume_from_checkpoint = True if glob.glob(output_dir + '/checkpoint-*') else False

trainer.train(resume_from_checkpoint=resume_from_checkpoint)
~~~

### 効果検証結果

それぞれの技術を個々適用した場合と総て同時に適用した場合の、使用GPUメモリと５エポック学習時間の比較表は以下です。

||										|使用GPUメモリ		|5エポック学習時間		|
|--	|--									|:--:				|:--:					|
|削減技術適用前<BR>(基準)|					|12589MiB			|13時間49分				|
||	
|適用する削減技術|Gradient Checkpointing	|8311MiB<BR>(34%削減)	|16時間36分<BR>(1.20倍)		|
||***Gradient Accumulation***			|3503MiB<BR>(72%削減)	|16時間28分<BR>(1.19倍)		|
||Mixed Precision Training				|10119MiB<BR>(20%削減)	|14時間44分<BR>(1.07倍)		|
||Optimizer 8-bit Quantization			|12621MiB<BR>(0%増加)	|14時間23分<BR>(1.04倍)		|
||
|***削減技術全適用***|					|***2873MiB<BR>(77%削減)***|***19時間48分<BR>(1.43倍)***|
|(参考) Optimizer<BR>8-bit Quantization<BR>以外総て適用||3183MiB<BR>(75%削減)|20時間55分<BR>(1.51倍)	|

GPUメモリ削減技術適用前と比較すると、削減技術全適用の場合、77％メモリが削減されますが、学習時間も43%長くなります。

なお、Optimizer 8-bit Quantizationだけを適用した場合、GPUメモリ削減技術適用前よりもGPU使用メモリが若干多くなったため、本当に効果があるか確認するためにOptimizer 8-bit Quantization以外総て適用した場合も測定しました。全適用の場合よりもGPU使用メモリは多く、また学習時間も長くなりますので、Optimizer 8-bit Quantizationの効果はあるようです。

GPUメモリ削減技術適用前、削減技術全適用時、Optimizer 8-bit Quantization以外総て適用時、それぞれの学習曲線は以下です。全適用では初期段階で学習が速まる傾向を見せましたが、その後はどれも同じような傾向になりました。これにより、GPUメモリを削減しても学習精度へは影響がなさそうなことがわかりました。

![キャプチャ.JPG](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/629790/3d27bb8b-9e4c-7e74-11af-c4db5a2620a3.jpeg)


# 10億超パラメータのRoBERTaモデルの学習

実際に、GPUメモリ削減技術が効果があることを確認できたので、いよいよ10億を超えるパラメータの[RoBERTaモデル][RoBERTa実装説明]を定義し、学習します。ここでも効果検証で使用したのと同じ[Notebook]を使用します。


### Model Config

[Notebook]「We'll define the following config for the model」の`config`を以下に変更します。


~~~python
config = RobertaConfig(
    vocab_size=52_000,
    max_position_embeddings=514,
    num_attention_heads=25,
    num_hidden_layers=48,
    type_vocab_size=1,
    hidden_size=1500,
    intermediate_size=4096,
)
~~~

この設定は[RoBERTaモデル][RoBERTa実装説明]のパラメータ数を10億超とするために[事前学習モデルの情報][Pretrained models]を参考にしました。

|モデル名		|パラメータ数|hidden<BR>層数|hidden<BR>サイズ|heads<BR>数|備考		|
|--				|:--:	|:--:|:--:|:--:	|--										|
|[roberta-base]	|125M	|12	|768	|12	|RoBERTa using the BERT-base<BR>architecture	|
|[roberta-large]|355M	|24	|1024	|16	|RoBERTa using the BERT-base<BR>architecture	|
|gpt2			|117M	|12	|768	|12	|OpenAI GPT-2 English model						|
|gpt2-medium	|345M	|24	|1024	|16	|OpenAI’s Medium-sized<BR>GPT-2 English model	|
|gpt2-large		|774M	|36	|1280	|20	|OpenAI’s Large-sized<BR>GPT-2 English model	|
|gpt2-xl		|1558M	|48	|1600	|25	|OpenAI’s XL-sized<BR>GPT-2 English model		|

`model.num_parameters()`によりこのモデルのパラメータ数が***1,103,750,608*** (***1.1B***)であることがわかります。


### Training Config

GPUメモリ削減技術の設定は[GPUメモリ削減技術の設定](#gpuメモリ削減技術の設定)で述べたコンフィグを用います。[Google Colab Pro]のセッションが学習途中で終了しても、次のセッションに効率的に引き継げるように`save_steps=500`でこまめにチェックポイントをセーブします。また、各チェックポイントで生成されるファイルサイズが大きいため、ディスクに負担をかけないように`save_total_limit=1`でチェックポイントファイルを常時１つに絞ります。

この設定を盛り込み、[Notebook]「Finally, we are all set to initialize our Trainer」のセルを以下に書き換えて実行します。

~~~python
from transformers import Trainer, TrainingArguments
from transformers.training_args import OptimizerNames

output_dir="/gpt2/EsperBERTo-1B"

# 小さなミニバッチサイズ4、疑似的ミニバッチサイズ64
training_args_4 = TrainingArguments(
    output_dir=output_dir,
    overwrite_output_dir=True,
    num_train_epochs=1,
    #per_gpu_train_batch_size=64,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=16,
    save_steps=500,
    save_total_limit=1,
    prediction_loss_only=True,
    fp16=True, # FP16の場合
    gradient_checkpointing=True,
    optim=OptimizerNames.ADAMW_BNB,
)

trainer = Trainer(
    model=model,
    args=training_args_4,
    data_collator=data_collator,
    train_dataset=dataset,
)
~~~

### Training実行

いよいよ学習です。一回エポックを実行するのに50時間程かかりますが、[Google Colab Pro]はセッションが大体一日以内に終了してしまうので、継続実行時には、チェックポイントから実行できるように、[Notebook]「Start training」を以下に変更して実行します。

~~~python
import os, glob

resume_from_checkpoint = True if glob.glob(output_dir + '/checkpoint-*') else False

trainer.train(resume_from_checkpoint=resume_from_checkpoint)
~~~

これにより、以下のように学習が始まります。

![キャプチャ.JPG](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/629790/1e45767b-6bf1-dc13-25d9-bf2ca02e4908.jpeg)

以下はセッション終了後に再度セッションを立ち上げ、学習を継続した際のキャプチャです。

![キャプチャ2.JPG](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/629790/befe76bc-6447-de8a-f4a1-cc43a04f662f.jpeg)


1エポック終了時のキャプチャです。

![キャプチャ.JPG](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/629790/f0b9d0cc-3a57-fccd-6fe1-734c12901f37.jpeg)



### Training実行時にCUDA out of memoryがでた場合の対処

[Google Colab Pro]で実行しても上記設定の場合、CUDA out of memoryがでる場合があります。

一つの原因は、本設定が16GB GPUメモリを念頭にチューンしたことにあります。[Google Colab Pro]はリソース割り当てを保証しているわけではないため、16GB GPUメモリよりも少ないGPUを割り当てることがあります。そうすると本設定ではエラーになってしまいます。この場合、筆者はセッションを終了し再度起動して16GBのGPUメモリを搭載するP100が割り当たるのを待ちました。GPUメモリサイズは`nvidia-smi`コマンドの出力結果の黄色でハイライトした箇所でご確認ください。

![キャプチャ.JPG](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/629790/05f56a29-254f-da5b-7d56-aa0ac8cdb5fe.jpeg)

また、小さなミニバッチサイズ8、疑似的ミニバッチサイズ64とした場合は、GPUメモリを16GBぎりぎりまで使います。実行できる場合もありますが、まだ余剰があるのにGPUメモリを確保できない、という以下のエラーが出る場合もあります。

![キャプチャ.JPG](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/629790/7e1d9056-9c6e-83f2-48fc-b8d88f72fa51.jpeg)

これには、セルで以下を実行することで回避できる場合がありますが、実行が少し遅くなるようですので、筆者は小さなミニバッチサイズを4とする本記事で採用した[コンフィグ](#gpuメモリ削減技術の設定)を使うことをお勧めします。

~~~python
%env PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:100
~~~

# 最後に

パラメータが10億超あるモデルでも、GPUメモリ16GBの実行環境で学習を1エポックまわすことができました。GPUメモリが足りないときはこれまでもミニバッチサイズを小さくしていましたが、Gradient Accumulationで大きなミニバッチサイズを疑似して精度をだすために活用できそうです。また、Gradient CheckpointingはO(sqrt(n))のスケールでメモリ消費量を抑えられるので、84Mパラメータモデルの4倍のレイヤー数とした11億パラメータモデルでは、[効果検証結果](#効果検証結果)でのメモリ削減より大きな削減がなされたものと思われます。

今回ご紹介したGPUメモリ削減技術が巨大なネットワークに対しても有効だということが分かりましたので、今後は[Google Colab Pro]でこれまで動かせなかった他のモデルに適用してどうなるか、見ていきたいと思っています。

最後までお読みいただきありがとうございました。

# 参考資料

- メモリ削減技術
  - [How to fine tune VERY large model if it doesn’t fit on your GPU]
  - Gradient Checkpointing
    - 論文：[Training Deep Nets with Sublinear Memory Cost](https://arxiv.org/abs/1604.06174v2)
    - Github：[cybertronai/gradient-checkpointing](https://github.com/cybertronai/gradient-checkpointing)
  - Gradient Accumulation
    - [How to Break GPU Memory Boundaries Even with Large Batch Sizes](https://towardsdatascience.com/how-to-break-gpu-memory-boundaries-even-with-large-batch-sizes-7a9c27a400ce)
    - [Gradient Accumulation in PyTorch]
  - Mixed Precision Training
    - 論文：[Mixed Precision Training](https://arxiv.org/abs/1710.03740)
    - [NVIDIA Train With Mixed Precision]
  - Optimizer 8-bit Quantization
    - 論文：[8-bit Optimizers via Block-wise Quantization](https://arxiv.org/abs/2110.02861)
    - Github：[facebookresearch/bitsandbytes](https://github.com/facebookresearch/bitsandbytes)
- RoBERTaモデル
  - 論文：[RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://arxiv.org/abs/1907.11692)
  - [Hugging FaceでのRoBERTa実装説明][RoBERTa実装説明]
  - Github：[huggingface/transformers][RoBERTa Github]
- RoBERTaモデルの学習方法
  - Blog：[How to train a new language model from scratch using Transformers and Tokenizers](https://huggingface.co/blog/how-to-train)
  - Notebook：[01_how-to-train.ipynb][Notebook]



[Google Colab Pro]: https://colab.research.google.com/signup

[How to fine tune VERY large model if it doesn’t fit on your GPU]: https://bestasoff.medium.com/how-to-fine-tune-very-large-model-if-it-doesnt-fit-on-your-gpu-3561e50859af

[transformers]: https://github.com/huggingface/transformers
[RoBERTa実装説明]: https://huggingface.co/docs/transformers/model_doc/roberta
[RoBERTa Github]: https://github.com/huggingface/transformers
[RoBERTa論文]: https://arxiv.org/abs/1907.11692
[RoBERTa-base]: https://dl.fbaipublicfiles.com/fairseq/models/roberta.base.tar.gz
[RoBERTa-large]: https://dl.fbaipublicfiles.com/fairseq/models/roberta.large.tar.gz
[Gradient Checkpointing]: https://arxiv.org/abs/1604.06174v2
[Gradient Checkpointing github]: https://github.com/cybertronai/gradient-checkpointing
[Gradient Accumulation in PyTorch]: https://kozodoi.me/python/deep%20learning/pytorch/tutorial/2021/02/19/gradient-accumulation.htm
[Mixed Precision Training]: https://arxiv.org/pdf/1710.03740v3.pdf
[NVIDIA Mixed Precision]: https://docs.nvidia.com/deeplearning/performance/mixed-precision-training/index.html
[NVIDIA Train With Mixed Precision]: https://docs.nvidia.com/deeplearning/performance/mixed-precision-training/index.html
[8-BIT OPTIMIZERS VIA BLOCK-WISE QUANTIZATION]: https://arxiv.org/pdf/2110.02861.pdf
[GLUE]: https://arxiv.org/abs/1804.07461
[How to train a new language model from scratch using Transformers and Tokenizers]: https://huggingface.co/blog/how-to-train
[Open in Colab]: https://colab.research.google.com/github/huggingface/blog/blob/main/notebooks/01_how_to_train.ipynb
[Notebook]: https://colab.research.google.com/github/huggingface/blog/blob/main/notebooks/01_how_to_train.ipynb

[Pretrained models]:https://huggingface.co/transformers/v2.2.0/pretrained_models.html