GaLore: コンシューマ向けハードウェアでの大規模モデル学習の進化

GaLoreは、NVIDIA RTX 4090 のようなコンシューマ向けハードウェア上で、最大70億パラメータの大規模言語モデル（LLM）の学習を可能にします。これは、オプティマイザ状態のメモリフットプリントを大幅に削減することで実現します。この AI 研究の民主化により、実務者は高価な産業用計算リソースを必要とせずに大規模モデルを学習できるようになります。

低ランク勾配射影によるメモリ効率化

GaLoreは、勾配をオプティマイザで処理する前に低次元の部分空間へ射影することでメモリ使用量を削減します。この手法は、深層ニューラルネットワークにおける勾配の固有の低ランク構造を利用し、学習中に保存・操作するデータ量を最小化します。

Adam のような適応的最適化アルゴリズムでは、オプティマイザ状態がメモリフットプリントの大部分を占めます。この射影を適用することで、GaLoreは学習中のオプティマイザ状態の保存に必要なメモリを82.5%以上削減したと報告されています。

動的部分空間スイッチング

全パラメータ学習の能力を維持し、モデルがパラメータ空間の限られた領域に閉じ込められるのを防ぐため、GaLoreは動的部分空間スイッチング機構を利用します。この機構により、モデルは学習過程で異なる低ランク部分空間を巡回できます。

これらのスイッチの頻度は、勾配の変化する低ランク構造に適応しつつ、一貫した最適化軌道を保つように調整されます。これにより、メモリ効率と最適化性能のトレードオフを微細に制御できます。

8ビットオプティマイザとの統合

GaLoreと8ビット精度オプティマイザを組み合わせることで、オプティマイザ状態を量子化し、メモリ効率をさらに最大化します。この相乗効果により、同じハードウェア制約下でモデル精度や収束速度を損なうことなく、より大きなモデルやより大きなバッチサイズでの学習が可能になります。

GaLore を用いた 8 ビット最適化のアルゴリズムプロセス

Gradient Projection: フル精度の勾配を射影行列で低ランク部分空間に射影し、8ビット形式に量子化します。
Quantization: 射影された勾配、モデルの重み、オプティマイザ状態（例：Adam の移動平均）を 32 ビット浮動小数点から 8 ビット整数表現に量子化します。
Optimizer Update: 8 ビットに量子化された勾配を更新します。これには、勾配を浮動小数点にデ量子化し、更新規則を適用し、更新されたオプティマイザ状態を再び 8 ビットに量子化する工程が含まれます。
De-quantization and Weight Update: 重みを処理のために浮動小数点にデ量子化します。その後、GaLore は最終的な射影を行い、デ量子化された低ランク更新を元のパラメータ空間にマッピングしてから重み更新を適用します。

Hugging Face Transformers での実装

GaLore は Hugging Face の transformers ライブラリ（バージョン 4.39.0 以上）および galore-torch ライブラリに統合されています。ユーザーは TrainingArguments でオプティマイザを galore_adamw、galore_adamw_8bit、galore_adafactor などのオプションで指定し、optim_target_modules で対象モジュールを定義することで GaLore を実装できます。

レイヤー単位の更新

メモリフットプリントをさらに削減するため、GaLore はレイヤー単位の重み更新をサポートします。逆伝播後に全レイヤーを一括で更新する代わりに、PyTorch のポストアキュムレーションフックを用いてオプティマイザがレイヤーごとに重みを更新します。この機能はオプティマイザ名に _layerwise を付加することで有効になります（例：galore_adamw_layerwise）。