Mini-batch Coresets for Memory-efficient Training of Large Language Models

要約

より大きなミニバッチを使用してトレーニングすると、収束率が向上し、優れたパフォーマンスが得られます。
ただし、大規模言語モデル (LLM) では、大規模な GPU メモリ要件が原因で、大規模なミニバッチを使用したトレーニングが法外になります。
この問題に対処するための効果的なアプローチは、より大きなミニバッチの勾配に厳密に一致する小さなミニバッチ コアセットを見つけることです。
ただし、言語データ内のソースの非常に不均衡な性質、Adam オプティマイザーの使用、および LLM の勾配次元が非常に大きいため、このアプローチは LLM では実行不可能で効果的ではありません。
この研究では、トレーニング LLM (CoLM) 用のコアセットを提案することで、上記の課題に対処します。
まず、勾配マッチングによって見つかったミニバッチ コアセットには、小さなソースの代表的な例が含まれていないため、最適なパフォーマンスを得るには、ミニバッチ コアセットに小さなソースのすべての例を含めることが重要であることを示します。
次に、履歴指数によって勾配を正規化し、Adam でトレーニングするためのミニバッチ コアセットを見つけます。
最後に、0 次法を活用して最後の V 射影行列の滑らかな勾配を見つけ、それをスパース化して、正規化された勾配の大きさが最大の次元を維持します。
CoLM を適用して、MathInstruct および SuperGLUE ベンチマークの LoRA を使用して Phi-2、Phi-3、および Zephyr を微調整します。
注目すべきことに、CoLM は微調整に必要なメモリを 2 分の 1 に削減し、さらに 4 倍大きいミニバッチを使用したトレーニングよりも優れたパフォーマンスを発揮します。
特に、CoLM は、LoRA などの既存のメモリ効率の高いトレーニング手法と簡単に積み重ねることができます。

要約(オリジナル)

Training with larger mini-batches improves the convergence rate and can yield superior performance. However, training with large mini-batches becomes prohibitive for Large Language Models (LLMs), due to the large GPU memory requirement. To address this problem, an effective approach is finding small mini-batch coresets that closely match the gradient of larger mini-batches. However, this approach becomes infeasible and ineffective for LLMs, due to the highly imbalanced nature of the sources in language data, use of the Adam optimizer, and the very large gradient dimensionality of LLMs. In this work, we address the above challenges by proposing Coresets for Training LLMs (CoLM). First, we show that mini-batch coresets found by gradient matching do not contain representative examples of the small sources w.h.p., and thus including all examples of the small sources in the mini-batch coresets is crucial for optimal performance. Second, we normalize the gradients by their historical exponential to find mini-batch coresets for training with Adam. Finally, we leverage zeroth-order methods to find smooth gradient of the last V -projection matrix and sparsify it to keep the dimensions with the largest normalized gradient magnitude. We apply CoLM to fine-tuning Phi-2, Phi-3, and Zephyr with LoRA on MathInstruct and SuperGLUE benchmark. Remarkably, CoLM reduces the memory requirement of fine-tuning by 2x and even outperforms training with 4x larger mini-batches. Notably, CoLM easily stack with existing memory-efficient training methods, such as LoRA.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google