Revisiting Zeroth-Order Optimization for Memory-Efficient LLM Fine-Tuning: A Benchmark

要約

自然言語処理 (NLP) の進化の状況では、SGD や Adam などの一次 (FO) オプティマイザーを使用して、事前トレーニング済みの大規模言語モデル (LLM) を微調整することが標準になっています。
しかし、LLM のサイズが大きくなるにつれて、FO 勾配計算のためのバックプロパゲーション (BP) による相当なメモリ オーバーヘッドが大きな課題となります。
この問題に対処することは、特にメモリ効率が最優先されるオンデバイス トレーニングのようなアプリケーションでは重要です。
この論文では、MeZO によって導入された初期コンセプトに基づいて、LLM 微調整中のメモリ コストを削減するソリューションとして、BP フリーのゼロ次 (ZO) 最適化への移行を提案します。
従来の ZO-SGD 手法とは異なり、私たちの研究では、5 つの LLM ファミリ (Roberta、OPT、LLaMA、Vicuna、Mistral) にわたる包括的な史上初のベンチマーク調査を通じて、調査をより広範囲の ZO 最適化手法に拡張しています。
3 つのタスクの複雑さと 5 つの微調整スキーム。
私たちの研究は、これまで見落とされていた最適化原理を明らかにし、タスクの調整の重要性、順勾配法の役割、アルゴリズムの複雑さとパフォーマンスの微調整のバランスを強調しています。
さらに、ブロックごとの降下、ハイブリッド トレーニング、勾配スパース性など、ZO 最適化に対する新しい機能強化を導入します。
私たちの研究は、さらなるメモリ効率の高い LLM 微調整を実現するための有望な方向性を提供します。
すべての実験を再現するコードは https://github.com/ZO-Bench/ZO-LLM にあります。

要約(オリジナル)

In the evolving landscape of natural language processing (NLP), fine-tuning pre-trained Large Language Models (LLMs) with first-order (FO) optimizers like SGD and Adam has become standard. Yet, as LLMs grow {in size}, the substantial memory overhead from back-propagation (BP) for FO gradient computation presents a significant challenge. Addressing this issue is crucial, especially for applications like on-device training where memory efficiency is paramount. This paper proposes a shift towards BP-free, zeroth-order (ZO) optimization as a solution for reducing memory costs during LLM fine-tuning, building on the initial concept introduced by MeZO. Unlike traditional ZO-SGD methods, our work expands the exploration to a wider array of ZO optimization techniques, through a comprehensive, first-of-its-kind benchmarking study across five LLM families (Roberta, OPT, LLaMA, Vicuna, Mistral), three task complexities, and five fine-tuning schemes. Our study unveils previously overlooked optimization principles, highlighting the importance of task alignment, the role of the forward gradient method, and the balance between algorithm complexity and fine-tuning performance. We further introduce novel enhancements to ZO optimization, including block-wise descent, hybrid training, and gradient sparsity. Our study offers a promising direction for achieving further memory-efficient LLM fine-tuning. Codes to reproduce all our experiments are at https://github.com/ZO-Bench/ZO-LLM .

arxiv情報

提供元, 利用サービス

arxiv.jp, Google