FZOO: Fast Zeroth-Order Optimizer for Fine-Tuning Large Language Models towards Adam-Scale Speed

要約

大型言語モデル（LLMS）の微調整は、GPUメモリボトルネックに直面することがよくあります。Adamのような1次オプティマイザーの後方パスは、メモリの使用量を推論レベルの10倍以上に増やします（たとえば、OPT-30Bで633 GB）。
Zeroth-Order（ZO）オプティマイザーは、前方パスからのみ勾配を推定することでこのコストを回避しますが、Mezoのような既存の方法は通常、収束するためにさらに多くのステップを必要とします。
ZOの速度とメモリ間のこのトレードオフは、基本的に改善できますか？
正規化されたSGDは、Adamよりもメモリ効率が大きい強力な経験的パフォーマンスを示しています。
これに照らして、Adamスケールの速度に向けて速いゼロオーダーオプティマイザーであるFZOOを紹介します。
FZOOは、バッチ損失の標準偏差に基づいてステップサイズを適応するバッチ片側の推定値を使用することにより、収束に必要な合計フォワードパスを減らします。
また、CUDAの並列処理と組み合わせたRademacherランダムベクトル摂動を使用して、バッチごとの計算を加速します。
Roberta-Large、OPT（350M-66B）、PHI-2、およびLLAMA3を含む多様なモデルでの広範な実験は、11のタスクを介してFZOOの有効性を検証します。
平均して、FZOOはメゾを精度で3％上回り、3倍の前方パスを必要とします。
Roberta-Largeの場合、FZOOは精度が5.6％の平均改善と、Mezoと比較して前方パスの18倍の減少を達成し、Adamに匹敵する収束速度を達成します。
また、正規化されたSGD更新ルールとその収束保証とFZOOの正式な同等性を証明する理論分析も提供します。
FZOOは、PEFTテクニックにスムーズに統合され、さらに大きなメモリの節約を可能にします。
全体として、我々の結果は、シングルGPU、高速、フルパラメーターの微調整の実用的であり、メモリ効率の高いトレーニングに関する将来の作業を指摘しています。

要約(オリジナル)

Fine-tuning large language models (LLMs) often faces GPU memory bottlenecks: the backward pass of first-order optimizers like Adam increases memory usage to more than 10 times the inference level (e.g., 633 GB for OPT-30B). Zeroth-order (ZO) optimizers avoid this cost by estimating gradients only from forward passes, yet existing methods like MeZO usually require many more steps to converge. Can this trade-off between speed and memory in ZO be fundamentally improved? Normalized-SGD demonstrates strong empirical performance with greater memory efficiency than Adam. In light of this, we introduce FZOO, a Fast Zeroth-Order Optimizer toward Adam-Scale Speed. FZOO reduces the total forward passes needed for convergence by employing batched one-sided estimates that adapt step sizes based on the standard deviation of batch losses. It also accelerates per-batch computation through the use of Rademacher random vector perturbations coupled with CUDA’s parallel processing. Extensive experiments on diverse models, including RoBERTa-large, OPT (350M-66B), Phi-2, and Llama3, across 11 tasks validate FZOO’s effectiveness. On average, FZOO outperforms MeZO by 3 percent in accuracy while requiring 3 times fewer forward passes. For RoBERTa-large, FZOO achieves average improvements of 5.6 percent in accuracy and an 18 times reduction in forward passes compared to MeZO, achieving convergence speeds comparable to Adam. We also provide theoretical analysis proving FZOO’s formal equivalence to a normalized-SGD update rule and its convergence guarantees. FZOO integrates smoothly into PEFT techniques, enabling even larger memory savings. Overall, our results make single-GPU, high-speed, full-parameter fine-tuning practical and point toward future work on memory-efficient pre-training.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google