Fine-tuning Quantized Neural Networks with Zeroth-order Optimization

要約

大規模な言語モデルのサイズが指数関数的に成長するにつれて、GPUメモリは、これらのモデルをダウンストリームタスクに適応させるためのボトルネックになりました。
この論文では、統一されたフレームワーク内でモデルの重み、勾配、およびオプティマイザー状態に関するメモリ使用量を最小限に抑えることにより、メモリ効率の高いトレーニングの限界を推進することを目指しています。
私たちのアイデアは、Zeroth-Orderの最適化を使用して勾配とオプティマイザーの両方の状態を排除することです。これは、前方パス中に摂動することで勾配に近いため、勾配方向を特定します。
重量のメモリ使用量を最小限に抑えるために、モデルの量子化、たとえばBFLOAT16からINT4への変換を採用しています。
ただし、ゼロオーダーの最適化を量子化された重みに直接適用することは、離散重みと連続勾配の間の精度のギャップのために実行不可能であり、それ以外の場合は定量化と再定量化が必要です。
この課題を克服するために、量子化されたZeroth-Order Optimization（QZO）を提案します。これは、勾配推定のための連続量子化スケールを妨害し、方向性微分クリッピング方法を使用してトレーニングを安定させる新しいアプローチを提案します。
QZOは、スカラーベースとコードブックベースのトレーニング後の量子化方法の両方に直交しています。
BFLOAT16のフルパラメーターの微調整と比較して、QZOは4ビットLLMSで総メモリコストを18ドル以上削減でき、1つの24GB GPU内で微調整されたLLAMA-2-13Bと安定した拡散3.5を有効にします。

要約(オリジナル)

As the size of large language models grows exponentially, GPU memory has become a bottleneck for adapting these models to downstream tasks. In this paper, we aim to push the limits of memory-efficient training by minimizing memory usage on model weights, gradients, and optimizer states, within a unified framework. Our idea is to eliminate both gradients and optimizer states using zeroth-order optimization, which approximates gradients by perturbing weights during forward passes to identify gradient directions. To minimize memory usage on weights, we employ model quantization, e.g., converting from bfloat16 to int4. However, directly applying zeroth-order optimization to quantized weights is infeasible due to the precision gap between discrete weights and continuous gradients, which would otherwise require de-quantization and re-quantization. To overcome this challenge, we propose Quantized Zeroth-order Optimization (QZO), a novel approach that perturbs the continuous quantization scale for gradient estimation and uses a directional derivative clipping method to stabilize training. QZO is orthogonal to both scalar-based and codebook-based post-training quantization methods. Compared to full-parameter fine-tuning in bfloat16, QZO can reduce the total memory cost by more than 18$\times$ for 4-bit LLMs, and enables fine-tuning Llama-2-13B and Stable Diffusion 3.5 Large within a single 24GB GPU.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google