LLM の学習を行う際は GPU メモリの消費量に悩まされることが多いです。
そういうときにどのように対応すればよいのか調べていたら huggingface が色々まとめてくれていたので、自分の勉強がてら要約していきます。

方法まとめ

学習速度とメモリ消費量の最適化に関係する方法には以下のようなものがあります。

方法／ツール	学習速度の向上	メモリ消費量の最適化
バッチサイズの変更	〇	〇
勾配蓄積 (Gradient accumulation)	×	〇
勾配チェックポイント (Gradient checkpointing)	×	〇
混合精度トレーニング (Mixed precision training)	〇	（×）
オプティマイザの変更	〇	〇
データプリロード (Data preloading)	〇	×
DeepSpeed Zero	×	〇
torch.compile	〇	×
Parameter-Efficient Fine Tuning (PEFT)	×	〇

バッチサイズの変更

バッチサイズと入出力ニューロン数は $2^N$ のサイズで使用することが推奨されます。しばしば8の倍数ですが、使用するハードウェアやモデルのdtypeに応じて、より高い数になることもあります。

Tensor Coreの要件は、dtypeとハードウェアに基づいて乗数を定義します。例えば、fp16データタイプでは8の倍数が推奨されますが、A100 GPUの場合は64の倍数を使用します。

適切な乗数が大幅な速度向上をもたらすことがあります。

個人的な話をすると、私はバッチサイズの選択は感覚的に行うことが多かったのですが、利用している GPU や dtype によって最適なサイズが決まってくるようです。

ただ、注意しておくと、これは速度とメモリ消費量の文脈の話で、学習後のモデルのパフォーマンスにとっての話ではないです。

勾配蓄積 (Gradient accumulation)

勾配蓄積法は、一度に全バッチの勾配を計算するのではなく、より小さなインクリメントで勾配を計算することを目指します。

このアプローチは、モデルを通じて前方および後方パスを実行することで小バッチの勾配を反復的に計算し、プロセス中に勾配を蓄積することを含みます。十分な数の勾配が蓄積されたら、モデルの最適化ステップが実行されます。

勾配蓄積を使用することで、GPUのメモリ容量によって課される制限を超えて実質的なバッチサイズを増加させることが可能になります。

勾配蓄積は、ミニバッチ全体の勾配を一度で計算するのではなく、複数回に分けて計算を行う方法です。
一度の計算ではメモリがあふれてしまう場合でも、分割して計算することによって計算が可能になります。

ただし、当然のことながら、勾配蓄積の分割を増やすほど、速度は遅くなります。

勾配蓄積によって導入される追加の前方および後方パスがトレーニングプロセスを遅くする可能性があることに注意が必要です。

GPUの使用率をできるだけ最大化することが推奨されますが、勾配蓄積ステップの数が多いとトレーニングの遅延が顕著になることがあります。

勾配蓄積は TrainingArguments で以下のように変数を定義するだけで使えます。

training_args = TrainingArguments(per_device_train_batch_size=1, gradient_accumulation_steps=4, **default_args)

勾配チェックポイント (Gradient checkpointing)

大規模なモデルでは、バッチサイズを1に設定し、勾配蓄積を使用してもなお、メモリの問題が発生することがあります。

前方パスからすべてのアクティベーションを保存し、後方パスで勾配を計算すると、大きなメモリオーバーヘッドが発生する可能性があります。

勾配チェックポイントは、これら二つのアプローチの妥協点を提供し、計算グラフ全体で戦略的に選択されたアクティベーションを保存するため、勾配の再計算が必要なアクティベーションの一部だけで済みます。

勾配チェックポイントは勾配計算の途中結果を全て保存する代わりに、必要なときに再計算することでメモリを節約する手法です。
勾配チェックポイントは TrainingArguments で以下のように変数を定義するだけで使えます。

training_args = TrainingArguments(
    per_device_train_batch_size=1, gradient_accumulation_steps=4, gradient_checkpointing=True, **default_args
)

混合精度トレーニング (Mixed precision training)

混合精度トレーニングは、特定の変数に対して低精度の数値形式を利用することで、モデルのトレーニングの計算効率を最適化する技術です。

伝統的に、ほとんどのモデルは32ビット浮動小数点精度（fp32またはfloat32）を使用して変数を表現および処理します。しかし、正確な結果を得るためにすべての変数がこの高精度レベルを必要とするわけではありません。

このアプローチでは、一部の計算が半精度で行われる一方で、一部は完全精度で行われるため、混合精度トレーニングと呼ばれます。

上記の説明の通りですが、混合精度トレーニングは fp32 ではなく、fp16 や bf16 等の浮動小数点で計算する手法です。

fp16

fp16 で計算する場合は以下のように fp16 フラグを True に設定します。

training_args = TrainingArguments(per_device_train_batch_size=4, fp16=True, **default_args)

bf16

bf16 で計算する場合は以下のように bf16 フラグを True に設定します。

training_args = TrainingArguments(bf16=True, **default_args)

tf32

他にも tf32 を利用する方法もあります。

Ampereハードウェアはtf32と呼ばれる魔法のようなデータタイプを使用します。

tf32サポートを有効にすることで最大3倍のスループット向上が得られる

ハードウェアの制限があるようですが、tf32 を使うと精度をそこまで落とさず計算速度が向上できるようです。

import torch
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

training_args = TrainingArguments(tf32=True, **default_args)

オプティマイザの変更

最も一般的なオプティマイザーはAdamまたはAdamW（ウェイト減衰付きAdam）ですが、GPUメモリの使用量を減らすために他のオプティマイザーを選択することも有効です。

optimizer には脳死で性能の良い Adam を利用する人も多いでしょうが、Adam はめちゃくちゃメモリを消費します。
そのため、optimizer を試行錯誤するのは効果的です。

具体的には以下のような選択肢があります。

adafactor
8-bit AdamW

adafactor

Adafactorは各重み行列の要素の移動平均を保存せず、行ごとおよび列ごとの移動平均の合計情報を保持することでメモリ使用量を大幅に削減します。ただし、Adamと比較して一部のケースで収束が遅くなる可能性があります。

adam のメモリ効率よい版です。

以下のように指定します。

training_args = TrainingArguments(per_device_train_batch_size=4, optim="adafactor", **default_args)

8-bit AdamW

このオプティマイザーは、オプティマイザーの状態を量子化することでメモリ使用量を大幅に削減します。3Bパラメータモデルの場合、24GBのGPUメモリを必要とする標準のAdamWオプティマイザーに対して、このオプティマイザーはたったの6GBで済みます。

AdamW の量子化版。量子化しているのでメモリ消費量が小さい。ただし精度は落ちる。

training_args = TrainingArguments(per_device_train_batch_size=4, optim="adamw_bnb_8bit", **default_args)

データプリロード (Data preloading)

優れたトレーニング速度に到達するための重要な要件の1つは、GPUが処理できる最大速度でGPUにデータを渡せることです。デフォルトでは、すべてがメインプロセスで行われるため、ディスクからデータを十分に速く読み取ることができず、ボトルネックとなり、GPUが十分に活用されない可能性があります。ボトルネックを減らすには、以下の引数を設定します

GPU を効率的に利用して学習速度を向上させるための手法です。以下のように実装します。

DataLoader(pin_memory=True, ...)
- データを CPU メモリにプリロードし、CPU から GPU メモリへの転送を大幅に高速化する
DataLoader(num_workers=4, ...)
- データをより速くプリロードするために、複数のワーカーを利用する
- GPU の利用率が低いなら、ワーカーの数を増やすと効果的

DeepSpeed Zero

DeepSpeedは、大規模なディープラーニングトレーニングの効率とスケーラビリティを向上させるために設計された機能と最適化を提供する、オープンソースのディープラーニング最適化ライブラリです。🤗 Transformersおよび🤗 Accelerateと統合されており、特に大規模モデルやメモリを多く必要とするシナリオでその強みを発揮します。

モデルが1つのGPUに収まらない場合や、小さなバッチサイズに収まらない場合は、DeepSpeed ZeRO + CPU Offloadや、はるかに大きなモデル用のNVMe Offloadを活用できます。

複数 GPU での学習をする際に DeepSpeed を使うとメモリ効率よく学習ができるようです。

torch.compile

PyTorch 2.0が導入したtorch.compileは、既存のPyTorchコードに手を加えることなく、モデルのコードを最適化する機能です。単一のコード行を追加するだけでモデルをコンパイルし、実行効率を向上させることができます。

どれくらい効果があるかわかりませんが、torch.compileでモデルのコードを最適化したら速度が向上するみたいですね。

以下のように実装します。

# trainer を利用しない場合
model = torch.compile(model)
# trainer を利用する場合
training_args = TrainingArguments(torch_compile=True, **default_args)

Parameter-Efficient Fine Tuning (PEFT)

PEFT（Parameter-Efficient Fine Tuning）メソッドは、微調整中に事前学習されたモデルのパラメータを凍結し、その上に少数の学習可能なパラメータ（アダプタ）を追加する。

LoRA のような Adapter を追加して学習したりする PEFT ですね。ここに関しては省略しますが、もし上記のようなことも試して学習できないようなら、Full-Parameter Fine-Tuning ではなく、PEFT の利用も考えないといけないですね。

まとめ

以上、LLM 学習時の速度やメモリに関する huggingface のドキュメントまとめでした。
知らないことも多く意外と勉強になりました。どれくらい効率的に学習できるのか個人的に色々試してみたいですね。

おわり