Adam-mini: Use Fewer Learning Rates To Gain More

要約

我々は、AdamWと同等かそれ以上の性能を、45%から50%少ないメモリフットプリントで達成するオプティマイザAdam-miniを提案する。Adam-miniは、Adamの学習率リソース（すなわち$1/sqrt{v}$）を削減することで、メモリを削減する。もし、(1)我々が提案するヘシアン構造に関する原理に従って、パラメータを注意深くブロックに分割する；(2)各パラメータブロックに、単一だが良い学習率を割り当てる。さらに、これらの各パラメータブロックに対して、それを探索するのに十分な資源があれば、アダムを凌駕する単一の高品質な学習率が存在することを見出す。そして、優れた学習率を見つけるためのコスト効率の良い一つの方法を提供し、Adam-miniを提案する。経験的に、Adam-miniは、事前学習、教師付き微調整、RLHFにおいて、125Mから7Bまでの様々なサイズの言語モデルで、AdamWと同等以上の性能を発揮することを検証する。Adam-miniのメモリフットプリントの削減は、GPUとCPU間の通信オーバーヘッドを軽減し、スループットを向上させます。例えば、Adam-miniは、$2times$ A800-80GBのGPUでLlama2-7Bを事前学習するとき、AdamWより49.6%高いスループットを達成し、事前学習にかかるウォールクロック時間を33%節約します。

要約(オリジナル)

We propose Adam-mini, an optimizer that achieves on-par or better performance than AdamW with 45% to 50% less memory footprint. Adam-mini reduces memory by cutting down the learning rate resources in Adam (i.e., $1/\sqrt{v}$). We find that $\geq$ 90% of these learning rates in $v$ could be harmlessly removed if we (1) carefully partition the parameters into blocks following our proposed principle on Hessian structure; (2) assign a single but good learning rate to each parameter block. We further find that, for each of these parameter blocks, there exists a single high-quality learning rate that can outperform Adam, provided that sufficient resources are available to search it out. We then provide one cost-effective way to find good learning rates and propose Adam-mini. Empirically, we verify that Adam-mini performs on par or better than AdamW on various language models sized from 125M to 7B for pre-training, supervised fine-tuning, and RLHF. The reduced memory footprint of Adam-mini also alleviates communication overheads among GPUs and CPUs, thereby increasing throughput. For instance, Adam-mini achieves 49.6% higher throughput than AdamW when pre-training Llama2-7B on $2\times$ A800-80GB GPUs, which saves 33% wall-clock time for pre-training.

arxiv情報

提供元, 利用サービス

arxiv.jp, DeepL