GaLore - 家庭用ハードウェアでの大規模モデルの学習

以下の記事が面白かったので、簡単にまとめました。

・GaLore: Advancing Large Model Training on Consumer-grade Hardware

1. GaLore

「GaLore」は、「NVIDIA RTX 4090」などの家庭用GPU上で、Llamaなどの最大7Bパラメータを持つモデルの学習を容易にします。これは、学習プロセス中のオプティマイザの状態と勾配に従来関連付けられていたメモリ要件を大幅に削減することによって実現されます。

2. オプティマイザ状態でのメモリ効率

オプティマイザ状態は、特にAdamのような適応最適化アルゴリズムでは、モデルの学習中のメモリフットプリントの重要な部分を占めます。「GaLore」は、オプティマイザによって処理される前に、勾配を低次元の部分空間に投影することでこの問題に対処します。これにより、これらの状態を保存するために必要なメモリが削減されるだけでなく、最適化プロセスの有効性も維持されます。

メモリの節約は大幅であり、著者らは「学習中にオプティマイザの状態を保存するためのメモリが 82.5% 以上削減された」と報告しており、同じメモリ制約内でより大きなモデルを学習したり、より大きなバッチサイズを使用したりすることが可能になります。8bit精度のオプティマイザと組み合わせると、これらの節約はさらに顕著になります。

3. 部分空間スイッチングと高度な投影技術

「GaLore」の有効性の重要な要素は、動的部分空間切り替えメカニズムです。これにより、モデルは学習プロセス全体を通じて、さまざまな低ランクの部分空間をナビゲートできるようになります。これにより、モデルがパラメータ空間の限られた部分に限定されなくなり、フルパラメータ学習の能力が維持されます。部分空間をいつどのように切り替えるかの決定は非常に重要であり、これらの切り替えの頻度は、一貫した最適化軌道の維持と、勾配の低ランク構造の進化する状況への適応との間のバランスになります。

勾配構造の変化に応じてこれらの予測を動的に調整する機能は、「GaLore」 の強力なツールであり、大規模なモデルの学習に固有のメモリ最適化のトレードオフをより細かく制御できるようになります。

4. GaLoreと8bitオプティマイザの組み合わせ

「GaLore」と「8bitオプティマイザ」を組み合わせることで、学習プロセスの整合性とパフォーマンスを維持しながらメモリ効率を最大化する相乗効果が得られます。8bitオプティマイザは、オプティマイザの状態を量子化することでメモリフットプリントを削減します。「GaLore」の投影メカニズムと組み合わせて使用​​すると、モデルの精度や収束速度に妥協のない、メモリ効率の高い学習体制が得られます。

この組み合わせは、家庭用ハードウェアで大規模なモデルを学習したり、メモリに制約のある環境にモデルを展開したりするなど、メモリが重大なボトルネックになっているシナリオで特に効果的です。これにより、同じハードウェア制約内でより複雑なモデルや大規模なデータセットを使用できるようになり、限られたリソースで達成できる限界が押し広げられます。

5. 実装の詳細

LLMを学習するために8bitオプティマイザを「GaLore」と統合するには、勾配、重み、およびオプティマイザの状態を8bit表現に量子化する必要があります。この量子化プロセスによりメモリフットプリントが大幅に削減され、同じメモリ制約内でより大きなモデルの学習やより大きなバッチサイズの使用が可能になります。この統合のアルゴリズムの詳細には、いくつかの重要なステップが含まれており、その一部はネイティブCUDA実装によって効率が大幅に向上します。「GaLore」は、これらの手法を行列の量子化および特殊なパラメータ化とさらに緊密に統合する新しい可能性を開き、メモリ使用量のさらなる削減につながる可能性があります。現在、bitsandbytesでこの方向性を模索しています。

6. GaLoreによる8bit最適化のアルゴリズムの概要

6-1. 勾配投影

「GaLore」は、投影行列を使用して、完全精度の勾配を低ランクの部分空間に投影します。このステップにより、勾配の次元が削減され、その後8bit形式に量子化されます。

6-2. 量子化

投影された勾配は、モデルの重みとオプティマイザの状態 (Adam の移動平均など) とともに、32bit浮動小数点表現から8bit整数表現に量子化されます。これには、浮動小数点値を8bit範囲にスケーリングし、最も近い整数に丸めることが含まれます。

6-3. オプティマイザー更新

 8bitの量子化された勾配を使用してモデルの重みを更新します。このステップには、勾配を浮動小数点形式に逆量子化し、オプティマイザの更新ルール (アダムのモーメント更新やパラメータ調整など) を適用し、更新されたオプティマイザの状態をストレージ用に8bitに量子化して戻すことが含まれます。

6-4. 逆量子化と重みの更新

8bit量子化された重みは、処理のために浮動小数点表現に逆量子化されますが、値の範囲が限られているため、量子化形式に固有の8bit精度は保持されます。PyTorchなどのフレームワークの標準操作は8bit整数をサポートしておらず、そのような整数の重みは勾配に対応できないため、この手順が必要です。このアプローチは本質的に精度を向上させるものではありませんが、現在の深層学習ライブラリの制約内で量子化された重みの実際の適用と勾配計算を容易にします。逆量子化後、重み更新を適用する前に、「GaLore」は逆量子化された低ランク更新を元の空間に投影するもう1つの投影を使用することに注意してください。

7. HuggingFace Transformers

「GaLore」を「HuggingFace Transformers」で使用するには、最新のtransformersとgalora-torchをインストールします。

pip install -U transformers galore-torch

以下は、「imdb」データセットで「Mistral-7B」を事前学習する例です。

import torch
import datasets
from transformers import TrainingArguments, AutoConfig, AutoTokenizer, AutoModelForCausalLM
import trl

train_dataset = datasets.load_dataset('imdb', split='train')

args = TrainingArguments(
    output_dir="./test-galore",
    max_steps=100,
    per_device_train_batch_size=2,
    optim="galore_adamw",
    optim_target_modules=["attn", "mlp"]
)

model_id = "mistralai/Mistral-7B-v0.1"

config = AutoConfig.from_pretrained(model_id)

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_config(config).to(0)

trainer = trl.SFTTrainer(
    model=model, 
    args=args,
    train_dataset=train_dataset,
    dataset_text_field='text',
    max_seq_length=512,
)

trainer.train()

TrainingArgumentsには、有効なoptim_target_modules (単一の文字列、正規表現、文字列や正規表現のリストをサポート) と、optimの場合は有効なGaLoreオプティマイザ(galore_adamw, galore_adamw_8bit, galore_adafactorなど)を渡します。

8. レイヤーごとの更新

言及すべきもう1つの重要な点は、レイヤーごとのオプティマイザ (つまり、一度に1レイヤーずつ重みを更新する) です。 通常、オプティマイザはバックプロパゲーション後にすべての層に対して単一の重み更新を実行します。 これは、重み勾配全体をメモリに保存することによって行われます。 レイヤーごとの重み更新を採用することで、学習中のメモリ使用量をさらに削減できます。 内部では、これはユーザーが更新したいレイヤー上のPyTorchポストアキュムレーションフックを使用して実装されています。

この機能を使用するには、オプティマイザー名に _layerwise を追加するだけです (例: galore_adamw_layerwise)。

OpenAI GPT-4V／ChatGPT／GPTs 人工知能プログラミング実践入門 | 株式会社ボーンデジタル