Gradient Multi-Normalization for Stateless and Scalable LLM Training

要約

大規模な言語モデル（LLMS）のトレーニングは、通常、Adam（Kingma＆BA、2015）のような適応オプティマイザーに依存しており、収束を加速するために追加の状態情報を保存しますが、重要なメモリオーバーヘッドが発生します。
Swan（Ma et al。、2024）などの最近の取り組みは、瞬間的な勾配に適用されるマルチステップ前処理手順を介してAdamに匹敵するパフォーマンスを達成しながら、オプティマイザー状態の必要性を排除することによりこれに対処しています。
白鳥の成功に動機付けられ、複数の規範に従って確率的勾配を正常化するステートレスオプティマザーを設計するための新しいフレームワークを紹介します。
これを達成するために、これらの規範w.r.tの勾配の正規化を実施するための簡単な交互スキームを提案します。
私たちの手順は、arbitrary意的な精度まで、問題の固定点を生み出すことができ、白鳥が慎重に選択された規範を備えた私たちのアプローチの特定の例であり、そのデザインをより深く理解することができることを示しています。
ただし、スワンの計算上の高価なホワイトニング/直交段階は、大規模なLMSの実用性を制限しています。
原則的な視点を使用して、より効率的でスケーラブルで実用的なステートレスオプティマイザーを開発します。
私たちのアルゴリズムは、白鳥の特性を緩和し、メモリ効率を保持しながら計算コストを大幅に削減し、大規模なモデルのトレーニングに適用できます。
最大10億パラメーターを備えたトレーニング前のLlamaモデルの実験は、Adamよりも3倍のスピードアップを示しており、メモリ要件が大幅に削減され、他のメモリ効率の高いベースラインよりも優れています。

要約(オリジナル)

Training large language models (LLMs) typically relies on adaptive optimizers like Adam (Kingma & Ba, 2015) which store additional state information to accelerate convergence but incur significant memory overhead. Recent efforts, such as SWAN (Ma et al., 2024) address this by eliminating the need for optimizer states while achieving performance comparable to Adam via a multi-step preprocessing procedure applied to instantaneous gradients. Motivated by the success of SWAN, we introduce a novel framework for designing stateless optimizers that normalizes stochastic gradients according to multiple norms. To achieve this, we propose a simple alternating scheme to enforce the normalization of gradients w.r.t these norms. We show that our procedure can produce, up to an arbitrary precision, a fixed-point of the problem, and that SWAN is a particular instance of our approach with carefully chosen norms, providing a deeper understanding of its design. However, SWAN’s computationally expensive whitening/orthogonalization step limit its practicality for large LMs. Using our principled perspective, we develop of a more efficient, scalable, and practical stateless optimizer. Our algorithm relaxes the properties of SWAN, significantly reducing its computational cost while retaining its memory efficiency, making it applicable to training large-scale models. Experiments on pre-training LLaMA models with up to 1 billion parameters demonstrate a 3X speedup over Adam with significantly reduced memory requirements, outperforming other memory-efficient baselines.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google