APOLLO: SGD-like Memory, AdamW-level Performance

要約

大規模言語モデル (LLM) は、特に人気のある AdamW オプティマイザーを使用した場合、トレーニング中にメモリを大量に消費することで知られています。
このメモリ負荷により、より多くの、またはハイエンドの GPU を使用するか、バッチ サイズを減らす必要があり、トレーニングのスケーラビリティとスループットが制限されます。
これに対処するために、オプティマイザーのメモリ使用量を削減するために、さまざまなメモリ効率の高いオプティマイザが提案されています。
しかし、彼らは次のような重大な課題に直面しています。(i) コストのかかる SVD 運用への依存。
(ii) AdamW と比較してパフォーマンスの大幅なトレードオフ。
(iii) 競争力のあるパフォーマンスを維持するために、オプティマイザーのメモリ オーバーヘッドが依然として相当量あります。
この研究では、AdamW の学習率適応ルールが構造化された学習率の更新として効果的に粗くできることを確認しました。
この洞察に基づいて、メモリ効率の高い LLM 最適化のための近似勾配スケーリング (APOLLO) を提案します。これは、純粋なランダム射影に基づく補助的な低ランクのオプティマイザー状態を使用して学習率スケーリングを近似します。
この構造化された学習率更新ルールにより、APOLLO は同等の事前トレーニング パフォーマンスを提供しながら、さらなるメモリ削減に対する耐性が高くなります。
そのランク 1 のバリアントである APOLLO-Mini でさえ、SGD レベルのメモリ コストを備えた AdamW と比較して優れた事前トレーニング パフォーマンスを実現します。
広範な実験により、APOLLO シリーズは AdamW と同等かそれ以上のパフォーマンスを示しながら、AdamW の最適化状態をほぼ排除することでメモリの大幅な節約を達成できることが実証されました。
これらの節約により、システム レベルで大きなメリットがもたらされます。 (1) スループットの向上: 4 倍のバッチ サイズをサポートすることにより、AdamW と比較して 8xA100 ～ 80GB セットアップで 3 倍のスループット。
(2) モデルのスケーラビリティの向上: システム レベルの最適化を行わずに、A100-80GB GPU で単純な DDP を使用して LLaMA-13B を事前トレーニングします。
(3) ローエンド GPU フレンドリーな事前トレーニング: 重み量子化を使用して 12 GB 未満のメモリを使用して、単一の GPU で LLaMA-7B を事前トレーニングします。

要約(オリジナル)

Large language models (LLMs) are notoriously memory-intensive during training, particularly with the popular AdamW optimizer. This memory burden necessitates using more or higher-end GPUs or reducing batch sizes, limiting training scalability and throughput. To address this, various memory-efficient optimizers have been proposed to reduce optimizer memory usage. However, they face critical challenges: (i) reliance on costly SVD operations; (ii) significant performance trade-offs compared to AdamW; and (iii) still substantial optimizer memory overhead to maintain competitive performance. In this work, we identify that AdamW’s learning rate adaptation rule can be effectively coarsened as a structured learning rate update. Based on this insight, we propose Approximated Gradient Scaling for Memory-Efficient LLM Optimization (APOLLO), which approximates learning rate scaling using an auxiliary low-rank optimizer state based on pure random projection. This structured learning rate update rule makes APOLLO highly tolerant to further memory reductions while delivering comparable pre-training performance. Even its rank-1 variant, APOLLO-Mini, achieves superior pre-training performance compared to AdamW with SGD-level memory costs. Extensive experiments demonstrate that the APOLLO series performs on-par with or better than AdamW, while achieving greater memory savings by nearly eliminating the optimization states of AdamW. These savings provide significant system-level benefits: (1) Enhanced Throughput: 3x throughput on an 8xA100-80GB setup compared to AdamW by supporting 4x larger batch sizes. (2) Improved Model Scalability: Pre-training LLaMA-13B with naive DDP on A100-80GB GPUs without system-level optimizations. (3) Low-End GPU Friendly Pre-training: Pre-training LLaMA-7B on a single GPU using less than 12 GB of memory with weight quantization.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google