大規模モデルを単一GPUで効率的に学習する方法

2023年10月2日 18:24

以下の記事が面白かったので、かるくまとめました。

・Methods and tools for efficient training on a single GPU

1. LLMを単一GPUで効率的に学習する方法

大規模モデルの学習では、次の2つを考慮する必要があります。

・スループット・学習時間
・モデルのパフォーマンス

「スループット」 (サンプル / 秒) を最大化すると、学習コストの削減につながります。これは通常、GPUメモリを限界まで利用することで実現されます。必要なバッチサイズがメモリオーバーする場合は、「Gradient Accumulation」などの「メモリの最適化」が必要になります。

ただし、「推奨バッチサイズ」がメモリに収まる場合は、学習が遅くなる可能性があるため、「メモリの最適化」を適用する必要はありません。どのバッチサイズが最良の結果をもたらすかを決定し、それに応じてリソースを最適化する必要があります。

この記事で解説している学習方法は、次のとおりです。

2. バッチサイズの選択

効率的に学習するには、「適切なバッチサイズ」を特定することから始めます。「サイズ 2^N のバッチサイズ」と「入出力ニューロン数」が推奨されています。多くの場合8の倍数ですが、使用されているハードウェアとモデルの dtype によっては、これより大きくなる場合があります。

参考として、全結合層の入出力ニューロン数とバッチサイズに関する NVIDIAの推奨事項を確認してください。

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

パラメータが小さい場合は、Dimension Quantization Effects も考慮します。

3. Gradient Accumulation

「Gradient Accumulation」は、バッチ全体の勾配を一度に計算するのではなく、より小さな増分で計算する手法です。これにより、GPUメモリの制限を超えて有効なバッチサイズまで増やすことができます。

「Gradient Accumulation」の設定方法は、次のとおりです。
この例では、有効なバッチサイズは 4 になります。

training_args = TrainingArguments(
    per_device_train_batch_size=1,  # バッチサイズ
    gradient_accumulation_steps=4,  # Gradient Accumulation
    **default_args
)

「Gradient Accumulation」しない「per_device_train_batch_size=4」がGPUメモリの限界の時、バッチサイズ 64 で学習したい場合は、「per_device_train_batch_size=1」「gradient_accumulation_steps=64」ではなく、「per_device_train_batch_size=4」「gradient_accumulation_steps=16」を設定します。これによって、GPUメモリの限界まで利用しながら、有効なバッチサイズで処理できます。

4. Gadient Checkpointing

バッチサイズを1にしても、GPUメモリがオーバーする場合は、「Gradient Accumulation」を使用してさらなるメモリ最適化が可能です。ただし、学習速度はが約20%遅くなります。

「Gadient Checkpointing」の設定方法は、次のとおりです。

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

5. Mixed Precision Training

「Mixed Precision Training」は、特定変数に低精度の数値型を利用することで、学習モデルの計算効率を最適化する手法です。

最も一般的な「Mixed Precision Training」は、fp16 (float16) を使用しますが、一部のGPUアーキテクチャでは、bf16 / tf32を使用します。

5-1. FP16

「Mixed Precision Training」の主な利点は、アクティベーションを半精度 (fp16) で保存することにあります。勾配も半精度で計算されますが、最適化ステップで完全精度に変換されるため、ここではメモリは節約されません。「Mixed Precision Training」では計算が高速になりますが、バッチサイズが小さい場合、より多くのGPU メモリが使用される可能性もあります。これは、モデルが16bitと32bitの両方の精度 (GPU 上の元のモデルの1.5 倍) でGPU上に存在するためです。

「FP16」設定方法は、次のとおりです。

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

5-2. BF16

Ampere以降の場合は、「Mixed Precision Training」の学習と評価に bf16 を使用できます。bf16はfp16よりも精度が劣りますが、ダイナミックレンジははるかに大きくなります。fp16で指定できる最大の数値は65535であり、それを超える数値はオーバーフローになります。

「TF32」設定方法は、次のとおりです。

TrainingArguments(
    bf16=True,  # BF16
    **default_args
)

5-3. TF32

Ampereは、tf32 と呼ばれるデータ型を使用します。fp32 (8bit) と同じ数値範囲を持っていますが、23bit精度の代わりに10bit (fp16と同じ) しかなく、合計で 19bitしか使用しません。これは、通常の fp32 学習コードや推論コードを使用でき、tf32 サポートを有効にすることで最大3倍のスループット向上が得られます。

「TF32」設定方法は、次のとおりです。

TrainingArguments(
    tf32=True,  # TF32
    **default_args
)

6. Flash Attention 2

「Flash Attention 2」を統合すると、学習のスループットを高速化できます。

詳しくは、「Flash Attention 2」を参照してください。

7. オプティマイザの選択

Transformerモデルの学習に使用される最も一般的なオプティマイザは、「Adam」「AdamW」(重み減衰を備えたAdam) です。「Adam」は、前の勾配の移動平均を保存することで良好な収束を実現します。ただし、モデルパラメータ数と同じくらいの追加のメモリを消費します。これを解決するために、他のオプティマイザを使用できます。

Trainerで利用できる主なオプティマイザは、次のとおりです。

・adamw_hf
・adamw_torch
・adamw_torch_fused
・adamw_apex_fused
・adamw_anyprecision
・adafactor
・adamw_bnb_8bit

「NVIDIA/apex」がインストールされている場合、「adamw_apex_fused」が全AdamWの中で最速です。

「t5-3b」の場合、標準の「adamw」は24GB、「adafactor」は 12GB、「adamw_bnb_8bit」は6GBほど、GPUメモリが必要になります。

7-1. Adafactor

「Adafactor」は、重み行列の各要素の移動平均を保存しません。代わりに、集約された情報 (行および列ごとの移動平均の合計) が保持されるため、フットプリントが大幅に削減されます。ただし、Adamと比較して、場合によっては収束が遅くなる場合があります。

「Adafactor」の設定方法は、次のとおりです。

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

7-2. 8-bit Adam

「8-bit Adam」は、オプティマイザの状態を集約する代わりに、完全な状態を保持し、それを量子化します。量子化とは、状態をより低い精度で保存し、最適化の場合にのみ逆量子化することを意味します。

「 8-bit Adam」の設定方法は、次のとおりです。

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

7-3. multi_tensor

「pytorch-nightly」から「torch.optim._multi_tensor」が導入されました。これにより、多数の小さな特徴テンソルがある状況でのオプティマイザが大幅に高速化されます。最終的にはデフォルトになるはずですが、早く試してみたい場合は、GitHubのIssueを参照してください。

8. データのプリロード

優れた学習速度を実現するための重要な要件の1つは、GPUが処理できる最大速度でGPUにデータを供給できることです。デフォルトでは、すべての処理がメインプロセスで行われるため、ディスクからデータを十分な速度で読み取ることができない可能性があります。

次の引数を指定することで、データのプリロードを行い、ボトルネックを軽減できます。

・pin_memory : CPU上の固定メモリにデータを確実にプリロードし、CPUからGPUへの転送を大幅に高速化
・num_workers : データをより速くプリロードするために複数のワーカーを生成

データのプリロードの設定方法は、次のとおりです。

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

training_args = TrainingArguments(
    dataloader_num_workers=4,
    **default_args
)

9. DeepSpeed

「DeepSpeed」は、TransformersおよびAccelerate と統合されたオープンソースの深層学習最適化ライブラリです。

モデルが単一GPUで小さなバッチサイズに適合する場合は、「DeepSpeed」を使用する必要はありません。速度が低下するだけです。小さなバッチサイズでも適合できない場合は、「DeepSpeed ZeRO + CPUオフロード」またはより大規模な「NVMeオフロード」を活用できます。

詳しくは、「DeepSpeed Zero Integration」を参照してください。

10. torch.compile

「PyTorch 2.0」には、「model = torch.compile(model)」という1行を追加するだけでコードを最適化できる新しいコンパイル関数が追加されました。

「TrainingArguments」での「torch.compile」の設定方法は、次のとおりです。

training_args = TrainingArguments(
    torch_compile=True,  # torch.compile
    **default_args
)

この記事が気に入ったらサポートをしてみませんか？