十八龍掌：トランスフォーマーのメモリ使用量を最適化するこのスキルの組み合わせは、収集する価値があります

ビジョントランスフォーマーや LLM などのディープラーニングモデルをトレーニングする場合、ピーク時のメモリ消費は一般的なボトルネックになります。この記事では、モデリングのパフォーマンスと予測精度を犠牲にすることなく、PyTorch のメモリ消費を約 1/20 に削減できる一連のテクニックを紹介します。

この記事の著者は、PyTorch の Torchvision ライブラリのビジュアル トランスフォーマーをベースに、約 100 行のコードからなるトレーニング スクリプトを作成しました。すべてのコード例は GitHub でご覧いただけます。

この記事で紹介するテクノロジーの名称は次のとおりです。

• ビジョントランスフォーマーの微調整
• 自動混合精度トレーニング
• 低精度トレーニング
• バッチサイズの縮小トレーニング
• 勾配累積とマイクロバッチ
• よりスリムなオプティマイザーを選択する
• ターゲットデバイス上でモデルをインスタンス化する
• 分散トレーニングとテンソル共有
• パラメータのアンインストール
• 上記の 9 つの方法を組み合わせたものが LLM に使用できる包括的な方法であり、これは 10 番目の方法とも呼ばれます。

これらの方法は互いに分離されており、一緒に使用できます。

この論文の実験で使用された ViT は ViT-L-16 モデルです。研究者たちは、上記の方法を一つずつ追加した後、BigBird-Roberta LLM をトレーニングしてテキスト分類タスクを実行できるようにします。これらの技術により、このようなモデルを消費者向けハードウェアでトレーニングすることが可能になります。

ビジョントランスフォーマーの微調整

実験の PyTorch コードを簡素化するために、この記事ではオープンソース ライブラリ Fabric を使用します。これにより、わずか 12 行のコードでさまざまな高度な PyTorch テクノロジ (自動混合精度トレーニング、マルチ GPU トレーニング、テンソル シャーディングなど) を適用できます。

ネイティブの PyTorch コードと Fabric を使用するために変更されたコードの違いは微妙で、次のコードに示すように、わずかな変更のみが加えられています。

前述のように、変更は小さいですが、既存のコードを再構築することなく PyTorch の高度な機能を使用できるのは便利です。

上の図を要約すると、通常の PyTorch コードを PyTorch + Fabric に変換するための 3 つの主な手順は次のようにまとめられます。

1. Fabric をインポートし、Fabric オブジェクトをインスタンス化します。
2. Fabric を使用して、モデル、オプティマイザー、データ ローダーを設定します。
3. 一般的に使用される loss.backward() の代わりに fabric.backward() を呼び出して損失関数を構築します。

通常の PyTorch と Fabric を使用した PyTorch のパフォーマンスとメモリ消費はほぼ同じです。

プレーン PyTorch (01_pytorch-vit.py):

 Time elapsed 17.94 min Memory used: 26.79 GB Test accuracy 95.85%

PyTorch と Fabric (01-2_pytorch-fabric.py)

 Time elapsed 17.88 min Memory used: 26.84 GB Test accuracy 96.06%

次のコードを使用することもできます。

 model = vit_l_16(weights=ViT_L_16_Weights.IMAGENET1K_V1)

次と置き換えます:

 model = vit_l_16(weights=None)

置き換え後、微調整の代わりに同じ ViT アーキテクチャが最初からトレーニングされ、予測精度は 96% 以上から約 60% に低下します。

自動混合精度

前のセクションでは、Fabric を使用して PyTorch コードを変更しました。これに基づいて、混合精度と分散トレーニングを使用するには、コードを 1 行変更するだけで済みます。

混合精度トレーニングの適用

混合精度トレーニングを使用するには、次のコード行に少し変更を加えるだけです。

 fabric = Fabric(accelerator="cuda", devices=1)

次と置き換えます:

 fabric = Fabric(accelerator="cuda", devices=1, precisinotallow="16-mixed")

その後、予測精度を犠牲にすることなく、メモリ消費量は以下に示すように 26.84 GB から 18.21 GB に削減されました。

01-2_pytoch-fabric.py と 02_mixed-precision.py の結果の比較

さらに、混合精度トレーニングでは、メモリ使用量が削減されるだけでなく、実行時間が 6 分の 1 (17.88 分から 3.45 分) に短縮されます。これは予想外のメリットです。

混合精度トレーニングとは何ですか?

混合精度トレーニングでは、精度が失われないように 16 ビットと 32 ビットの両方の精度を使用します。 16 ビット表現による勾配計算は、32 ビット形式よりもはるかに高速で、大量のメモリを節約します。この戦略は、メモリや計算上の制約がある場合に特に有益です。

「低」精度トレーニングではなく「混合」精度トレーニングと呼ばれる理由は、すべてのパラメーターと操作が 16 ビット浮動小数点数に変換されるわけではないためです。実際には、トレーニング中に 32 ビット操作と 16 ビット操作を切り替えます。

下の図に示すように、混合精度トレーニングは、計算を高速化するために重みを低い精度 (FP16 など) に変換する、勾配を計算する、数値の安定性のために勾配をより高い精度 (FP32) に戻す、スケーリングされた勾配で元の重みを更新する、といういくつかのステップに分けられます。

この方法により、トレーニングの有効性を確保しながら、ニューラル ネットワークの精度と安定性を維持できます。

興味のある読者は、この記事の著者による別の記事「混合精度テクノロジーを使用した大規模言語モデルの高速化」で、より詳しい基礎概念を知ることもできます。

記事アドレス: https://lightning.ai/pages/community/tutorial/accelerating-large-language-models-with-mixed-precision-techniques/

低精度トレーニング

さらに一歩進んで、混合精度ではなく「完全な」下位 16 ビット精度で実行してみることもできます。

次のコード行を置き換えます。

 fabric = Fabric(accelerator="cuda", precisinotallow="16-mixed")

置き換える

fabric = Fabric(accelerator="cuda", precisinotallow="16-true")

ただし、トレーニング中に NaN 値が生成されることに注意してください。

 Epoch: 0001/0001 | Batch 0000/0703 | Loss: 2.4105 Epoch: 0001/0001 | Batch 0300/0703 | Loss: nan Epoch: 0001/0001 | Batch 0600/0703 | Loss: nan ...

これは、通常の 16 ビット浮動小数点では -65504 から 65504 までの数値しか表現できないためです。

 In [1]: import torch In [2]: torch.finfo(torch.float16) Out[2]: finfo(resolutinotallow=0.001, min=-65504, max=65504, eps=0.000976562, smallest_normal=6.10352e-05, tiny=6.10352e-05, dtype=float16)

したがって、NaN 問題を回避するには、パラメータを "bf16 true" に変更できます。

 fabric = Fabric(accelerator="cuda", precisinotallow="bf16-true")

メモリ消費量をさらに 13.82 GB まで削減することも可能です (この場合も、精度は犠牲になりません)。

03_bfloat16.pyの結果を前のコードと比較する

Bfloat16とは何ですか?

「bf16 mixed」の「bf16」は、Brain Floating Point (bfloat16) の略です。 Google は、特に Tensor Processing Unit (TPU) における機械学習およびディープラーニング アプリケーション向けにこの形式を開発しました。従来の float16 形式と比較すると、Bfloat16 は精度を犠牲にしてダイナミック レンジを拡張します。

拡張されたダイナミック レンジにより、bfloat16 は非常に大きな数値と非常に小さな数値を表すことができるため、幅広い範囲の値に遭遇する可能性のあるディープラーニング アプリケーションに適したものになります。ただし、精度が低いと、一部の計算の精度に影響したり、場合によっては丸め誤差が発生したりする可能性があります。しかし、ほとんどのディープラーニング アプリケーションでは、この精度の低下はモデリングのパフォーマンスにほとんど影響しません。

bfloat16 はもともと TPU 用に開発されましたが、この形式は A100 Tensor Core GPU 以降の NVIDIA GPU でもサポートされています。

次のコードは、GPU が bfloat16 をサポートしているかどうかを確認できます。

 >>> import torch >>> torch.cuda.is_bf16_supported() True

バッチサイズを縮小する

多くの場合、バッチ サイズを縮小することは、メモリ消費を削減する効果的な方法です。ただし、トレーニングのダイナミクスが変化するため、予測パフォーマンスが低下する場合があります。

いずれにしても、バッチ サイズを縮小すると結果にどのような影響があるかを調べる価値はあります。パフォーマンスを犠牲にすることなくバッチ サイズを 16 に減らすことが可能であり、メモリ消費量を 5.69 GB に削減できることがわかります。

04_lower-batchsize.py を前のコードと比較します。

勾配累積とマイクロバッチ処理

勾配累積は、トレーニング中にバッチ サイズを仮想的に増やす方法であり、使用可能な GPU メモリが目的のバッチ サイズに対応するのに十分でない場合に役立ちます。この方法は実行時にのみ影響し、モデリングのパフォーマンスには影響しません。

勾配累積では、各バッチの直後にモデルの重みを更新するのではなく、より小さなバッチが計算され、複数の反復にわたって勾配が累積されます (通常は合計または平均化されます)。累積勾配が目標の「仮想」バッチ サイズに達すると、モデルの重みは累積勾配で更新されます。

勾配累積を実装するには、前方パスと後方パスに 2 つの小さな変更を加えるだけです。

05_gradient-acum.py のコード修正

この記事の著者による別の記事「勾配累積を使用して単一 GPU 上で LLM を微調整する」では、勾配累積の詳細がさらに詳しく紹介されています。

記事アドレス: https://lightning.ai/blog/gradient-accumulation/

有効なバッチ サイズは 16 で、累積ステップ数は 4 なので、実際のバッチ サイズは 4 になります (16/4=4 のため)。

05_gradient-acum.pyの結果

この手法の欠点は、実行時間が 3.96 分から 12.91 分に増加することです。

注目すべきは、バッチ サイズを 1 まで小さくすることができ、メモリ消費量がさらに 75% 削減されることです。

よりスリムなオプティマイザーを使用する

人気の Adam オプティマイザーには、実際には追加のパラメーターが付属しています。たとえば、Adam は各モデル パラメーターに対して 2 つの追加のオプティマイザー パラメーター (平均と分散) を提供します。

したがって、Adam を SGD のようなステートレス オプティマイザーと交換することで、パラメーターの数を 2/3 に削減できます。これは、ViT と LLM を使用するときに非常に重要です。

通常の SGD の欠点は収束性が低いことです。したがって、Adam を SGD と交換した後は、補償するためにコサイン減衰学習率スケジューラを導入する必要があります。

つまり、次のコードを追加することで

optimizer = torch.optim.Adam(model.parameters(), lr=5e-5)

次と置き換えます:

 optimizer = torch.optim.SGD(model.parameters(), lr=0.01) num_steps = NUM_EPOCHS * len(train_loader) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=num_steps)

この変更により、モデルは約 97% の分類精度を維持しながら、ピーク時のメモリ消費量を削減できるようになりました。

06_sgd-with-scheduler.py の結果

ターゲットデバイス上でモデルを作成する

PyTorch でモデルをインスタンス化する場合、通常は最初に CPU デバイス上でモデルを作成し、次にそれをターゲット デバイスに転送して、必要な精度に変換します。

 model = vit_l_16(weights=ViT_L_16_Weights.IMAGENET1K_V1) model.cuda().float16()

ただし、CPU 上で完全精度の中間モデルを生成するのは非効率的な方法です。したがって、Fabric の init_module コンテキストを使用して、ターゲット デバイス (GPU など) 上で必要な精度のモデルを直接作成できます。

 import lightning as L fabric = Fabric(accelerator="cuda", devices=1, precisinotallow="16-true") with fabric.init_module(): model = vit_l_16(weights=ViT_L_16_Weights.IMAGENET1K_V1)

この特定のケース (モデル) では、フォワード パス中のピーク メモリは、完全精度表現でのサイズよりも大きくなります。モデルの読み込み自体の fabric.init_module メソッドをベンチマークすると、次の結果が得られます。

• init_module なしの GPU ピーク メモリ: 1.24 GB (07_01_init-module.py)
• init_module を使用した GPU ピーク メモリ: 0.65 GB (07_03_init-module.py)

この場合、init_module によってモデルの読み込みに必要なピーク メモリが 50% 削減されることがわかります。

init_module の詳細については、大規模モデルの効率的な初期化に関するこの記事を参照してください。

記事アドレス: https://lightning.ai/pages/community/efficient-initialization-of-large-models/

分散トレーニングとテンソル共有

次の変更は、マルチ GPU トレーニングです。複数の GPU を利用できると、モデルのトレーニングが高速化されるため便利です。

ただし、この記事ではメモリの節約について説明します。したがって、データ並列処理とテンソル並列処理を利用して複数のデバイス間で大きな重み行列を共有する、完全共有データ並列処理 (FSDP) と呼ばれる、より高度な分散マルチ GPU 戦略が必要です。

しかし、上記のセクション 7 のコードにこのテクニックを追加する場合など、モデルがすでに小さい場合は、目に見える効果はほとんどありません。したがって、シャーディングの効果にのみ焦点を当てるには、セクション 1 のフル精度ベースラインと比較することができます。

次のコードを置き換えてください

fabric = Fabric(accelerator="cuda", devices=1)

次と置き換えます:

 auto_wrap_policy = partial( transformer_auto_wrap_policy, transformer_layer_cls={EncoderBlock}) strategy = FSDPStrategy( auto_wrap_policy=auto_wrap_policy, activation_checkpointing=EncoderBlock) fabric = Fabric(accelerator="cuda", devices=4, strategy=strategy)

08_fsdpと-01-2.pyの結果

手動で定義するだけでなく、次の方法を使用して、分割するレイヤーを自動的に決定することもできます。

 fabric = Fabric(accelerator="cuda", devices=4, strategy="fsdp")

データ並列性とテンソル並列性を理解する

データ並列処理では、ミニバッチをさらに分割し、モデルのコピーを各 GPU に配置する必要があります。複数の GPU が並列に動作するため、モデルのトレーニングを高速化できます。

仕組みは次のとおりです:

1. すべての GPU にわたって同じモデルを複製します。
2. 次に、各 GPU に入力データの異なるサブセット (異なるミニバッチ) が供給されます。
3. すべての GPU はモデルの順方向パスと逆方向パスを独立して実行し、独自のローカル勾配を計算します。
4. 次に、勾配が収集され、すべての GPU にわたって平均化されます。
5. 平均化された勾配は、モデルのパラメータを更新するために使用されます。

この方法の主な利点は速度です。各 GPU は他の GPU と同時に独自のミニバッチ データを処理するため、モデルはより短時間でより多くのデータでトレーニングできます。これにより、特に大規模なデータセットを扱う場合に、モデルのトレーニングに必要な時間を大幅に短縮できます。

ただし、データの並列処理にはいくつかの制限があります。各 GPU には、モデルとそのパラメータの完全なコピーが必要です。これにより、モデルが単一の GPU のメモリに収まる必要があるため、トレーニング可能なモデルのサイズが制限されますが、これは最新の ViT または LLM では実現可能ではありません。

データ並列処理とは異なり、テンソル並列処理ではモデル自体が GPU 全体に分割されます。データ並列処理では、各 GPU がモデル全体に​​対応する必要があり、大規模なモデルをトレーニングするときに制限となる可能性があります。ただし、テンソル並列処理を使用すると、モデルを分割して複数のデバイスに分散することで、単一の GPU では大きすぎる可能性のあるモデルをトレーニングできます。

具体的には、その原理は行列の乗算に似ています。モデルは行または列ごとに分解できます。簡単にするために、列ごとの分解を例にとると、次の図に示すように、大規模な行列乗算演算を個別の計算に分解し、それぞれを異なる GPU で実行することができます。その後、結果が連結されて元の結果が得られ、計算負荷が効果的に分散されます。

パラメータのアンインストール

前のセクションで説明したFSDP戦略に加えて、次のコードを追加することで、最適化パラメータをCPUにオフロードすることもできます。

 strategy = FSDPStrategy( auto_wrap_policy=auto_wrap_policy, activation_checkpointing=EncoderBlock, )

次と置き換えます:

 strategy = FSDPStrategy( auto_wrap_policy=auto_wrap_policy, activation_checkpointing=EncoderBlock, cpu_offload=True )

メモリ消費量が 6.59 GB から 6.03 GB に減少しました。

09_fsdp-cpu-offload-with-01-2.py の結果。

ちょっとした欠点としては、上映時間が 5.5 分から 8.3 分に増えたことです。

これまでの技を組み合わせると、十八龍鎮圧掌の最も強力な最後の掌が完成します。

前のセクションでは ViT の最適化に関する多くのトピックを取り上げましたが、これらの手法は LLM にも適用できます。

著者らは、LLaMA、Falcon、Pythia、およびその他の一般的なモデルをサポートする Lit LLaMA および Lit GPT リポジトリでこれらの手法の多くを使用しています。それでも、より一般的な例を作成するために、著者らは、IMDb の映画レビューの感情を分類するための、人気の HF トランスフォーマー ライブラリから LLM を微調整しました。

上記の手法を使用すると、3.99 Gb (bonus_bigbird-before.py) ではなく 1.15 Gb のメモリ (bonus_DistilBERT-after.py) のみを使用して DistilBERT 分類器をトレーニングすることが可能になります。さらに印象的なのは、これらの手法をトランスフォーマー ライブラリの BigBird モデルに適用することで、BigBird が消費するサイズがわずか 4.03 GB (bonus_BigBird-after.py) になることです。

 strategy = FSDPStrategy( cpu_offload=True ) fabric = Fabric( accelerator="cuda", devices=4, strategy=strategy, precision="bf16-true" ) with fabric.init_module(): model = AutoModelForSequenceClassification.from_pretrained( "google/bigbird-roberta-base", num_labels=2)

結論は

この記事では、PyTorch モデルのメモリ消費を削減する 9 つの手法を紹介します。これらの技術を ViT に適用すると、単一の GPU でメモリ消費量が 20 倍削減されます。ご覧のとおり、テンソルを GPU 間でシャーディングすると、メモリ消費量も削減できます。同じ最適化により、BigBird LLM はピーク時の GPU RAM を 4GB のみ使用してトレーニングすることも可能になります。

これらのテクニックはいずれもモデル固有のものではなく、どの PyTorch トレーニング スクリプトでも使用できます。オープンソースの Fabric ライブラリを使用すると、ほとんどの最適化を 1 行のコードで実装できます。