大規模言語モデルのトレーニングの新しい手法
新しい方法が、一般のハードウェアで大きな言語モデルをトレーニングする際のメモリ効率を向上させる。
― 1 分で読む
目次
大規模言語モデル(LLM)をトレーニングするには、パラメータのサイズやトレーニングに使うオプティマイザの状態のせいで、たくさんのメモリが必要なんだ。これらのモデルは、チャットボットや翻訳といったタスクでどんどん一般的になってきてる。性能は良いけど、高メモリの要求があるから、特に個人用PCの普通のGPUではトレーニングが難しいんだよね。
トレーニングにおけるメモリの課題
典型的なLLMをトレーニングする時、使用するメモリは何十億ものパラメータ、さらにはその勾配やオプティマイザの状態を収容する必要があるんだ。例えば、LLaMA 7Bのようなモデルをゼロからトレーニングするには、58GB以上のメモリが必要になったりする。このメモリは、パラメータ自体や、Adamのようなオプティマイザからのモーメンタム項、トレーニング中の中間活性化で消費される。
だから、大規模なモデルのトレーニングは、たった24GBのメモリしかないNVIDIA RTX 4090のような消費者向けハードウェアでは通常不可能なんだ。
メモリ使用量を減らすための現在のアプローチ
研究者たちはトレーニング中のメモリ使用量を減らすためのいくつかの方法を試してきた。一つの人気な方法は低ランク適応(LoRA)っていうやつ。LoRAは、モデルの各レイヤーに小さくてトレーニング可能なマトリックスを追加するんだ。これによって、トレーニングする必要のあるパラメータの数が減って、メモリ使用料が下がるってわけ。
でも、LoRAや似たような技術は、フルセットのパラメータでトレーニングするのと比べると、パフォーマンスがあまり良くないことが多いんだ。モデルへの更新の範囲を制限しちゃうから、トレーニングがあまり効果的じゃなくなる。
勾配低ランク射影
新しいアプローチ:既存の方法の制限に対処するために、勾配低ランク射影っていう新しい戦略が提案された。このアプローチは、標準的な低ランク適応法(LoRAのような)よりもメモリ効率が良く、フルパラメータのトレーニングを可能にするんだ。
核心のアイデアは、トレーニング中にパラメータの勾配が低ランクの構造を持つことを利用すること。勾配の低ランクな部分に焦点を当てて、重みのマトリックス自体を低ランクに制限するのではなく、かなりのメモリの節約ができるんだ。
勾配マトリックスを低ランク形式に射影する技術を使うことで、オプティマイザの状態に必要なメモリを劇的に減らせる。実際に、この方法はトレーニング中のメモリ使用量を約65.5%削減しつつ、フルパラメータでトレーニングしたモデルと同等のパフォーマンスを維持することが示されてる。
パフォーマンスと効率
実際、この新しい方法は事前トレーニングやファインチューニングの両方に効果的だったんだ。C4データセットでLLaMAを事前トレーニングするとき、勾配低ランク射影技術と8ビットオプティマイザを組み合わせることで、モデルはフルランクの技術と比較しても十分なパフォーマンスを発揮しつつ、かなり少ないメモリで済むようになった。
さらに、消費者向けGPUを使って7億パラメータのモデルをトレーニングすることも可能にしたし、モデルパラレルやオフロードのような複雑な戦略は必要なかったんだ。
既存の方法に対する利点
フルランクのトレーニングの効果に感じるLoRAに比べて、勾配低ランク射影は明らかな優位性を維持してる。最初にフルランクモデルのトレーニングを要求することなく、メモリ使用量を減らせるんだよ。これはReLoRAのような方法のよくあるハードルなんだ。
この新しい技法は、既存のオプティマイザにシンプルな調整を加えるだけで済んで、AdamWやAdafactorのような人気モデルで効果的なパフォーマンスを示してくれてる。この方法によって追加されるハイパーパラメータも最小限で、実装が簡単なんだ。
メモリ効率的な技術に関する関連作業
モデルのトレーニング中のメモリフットプリントを減らすために、いくつかの戦略が既に提案されてる。例えば、Adafactor法はオプティマイザのセカンドオーダー統計をファクタリングすることでメモリ使用量を減らす。その他にも、勾配統計を圧縮してメモリ効率を向上させる方法が探求されてる。
低ランク戦略に関しては、サブスペース学習のような技術が特定の低次元空間内での更新を維持することを目指してる。こういった概念は期待できる結果を示してきたけど、複雑なタスク、特にLLMのトレーニングのようなものでは性能を制限する課題に直面することが多いんだ。
新しい方法の実装
勾配低ランク射影は、勾配マトリックスを低ランク形式に変換する射影マトリックスを定期的に計算することによって機能する。これにより、オプティマイザの状態を保存するために必要なメモリが減るんだ。こうしたメモリの節約がある一方で、計算のオーバーヘッドは最小限に抑えられるから、既存のトレーニングフレームワークにスムーズに統合できる。
この戦略を実装する上での実際的な側面は、低ランク射影を更新する頻度を決定することにある。サブスペースに対する変更は、何百回かのイテレーションごとに行っても、計算コストは大きく増えないんだ。
事前トレーニングとファインチューニングの利点
この方法の利点は、単なる事前トレーニングを超えて広がってる。GLUEのようなタスクでモデルをファインチューニングするとき、勾配低ランク射影は従来の方法(LoRAのような)と比較しても優れた結果を出してる。さまざまなタスクにおいてうまく機能して、メモリ使用量を抑えつつ、同等かそれ以上のパフォーマンスを維持してるんだ。
実験結果
LLaMAモデル上での広範なテストによると、勾配低ランク射影は他の低ランクの方法よりも常に優れたパフォーマンスを発揮してる。さまざまなモデルサイズに適用した結果、バリデーションの困惑度が低下し、フルランクのトレーニングの過剰なメモリ要求なしに効果的なモデルをトレーニングできる能力を示してる。
例えば、厳密な試験を通じて、この方法が事前トレーニング後のモデルのファインチューニングを効率的に扱えることが示されていて、特定のアプリケーションにモデルを適応させるためのスムーズなアプローチを提供してる。
現実世界への影響
消費者向けハードウェアで大規模な言語モデルを効果的にトレーニングできることは、重要な意味を持ってる。先進的なAI技術へのアクセスを民主化するだけでなく、こうしたモデルのトレーニングに伴う環境への影響を減らすことにもつながる。パワフルじゃないマシンでトレーニングを行えるようにすることで、エネルギー消費を最小限に抑え、全体的なカーボンフットプリントを削減できるんだ。
結論
勾配低ランク射影の導入は、大規模言語モデルのトレーニングにおけるメモリの課題を解決するための一歩だ。勾配の低ランク特性に焦点を当てることで、この方法は効果的な学習を可能にしつつ、メモリ使用量を大幅に削減できる。その結果、研究者や開発者が限られたハードウェア上でより大きくて複雑なモデルをトレーニングできるようになり、さまざまな分野でのAI技術の進歩につながるんだ。
今後の方向性
今後の研究では、モデルのトレーニングにおけるメモリ効率をさらに向上させることが探求できる。これは、効率的なパラメータ化や量子化方法を通じて低ランク射影マトリックスを強化することを含むかもしれない。また、このアプローチをビジョンタスクや生成モデルに使う可能性もあるんだ。
要するに、勾配低ランク射影は大規模言語モデルのトレーニングのための実用的で効率的な手段を提供して、よりアクセスしやすく持続可能なAI環境に寄与することになる。
タイトル: GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection
概要: Training Large Language Models (LLMs) presents significant memory challenges, predominantly due to the growing size of weights and optimizer states. Common memory-reduction approaches, such as low-rank adaptation (LoRA), add a trainable low-rank matrix to the frozen pre-trained weight in each layer, reducing trainable parameters and optimizer states. However, such approaches typically underperform training with full-rank weights in both pre-training and fine-tuning stages since they limit the parameter search to a low-rank subspace and alter the training dynamics, and further, may require full-rank warm start. In this work, we propose Gradient Low-Rank Projection (GaLore), a training strategy that allows full-parameter learning but is more memory-efficient than common low-rank adaptation methods such as LoRA. Our approach reduces memory usage by up to 65.5% in optimizer states while maintaining both efficiency and performance for pre-training on LLaMA 1B and 7B architectures with C4 dataset with up to 19.7B tokens, and on fine-tuning RoBERTa on GLUE tasks. Our 8-bit GaLore further reduces optimizer memory by up to 82.5% and total training memory by 63.3%, compared to a BF16 baseline. Notably, we demonstrate, for the first time, the feasibility of pre-training a 7B model on consumer GPUs with 24GB memory (e.g., NVIDIA RTX 4090) without model parallel, checkpointing, or offloading strategies.
著者: Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, Yuandong Tian
最終更新: 2024-06-02 00:00:00
言語: English
ソースURL: https://arxiv.org/abs/2403.03507
ソースPDF: https://arxiv.org/pdf/2403.03507
ライセンス: https://creativecommons.org/licenses/by/4.0/
変更点: この要約はAIの助けを借りて作成されており、不正確な場合があります。正確な情報については、ここにリンクされている元のソース文書を参照してください。
オープンアクセスの相互運用性を利用させていただいた arxiv に感謝します。