arXiv:2502.06742cs.LGcs.AI2025-02被引 9

提出无需状态的高效优化器,训练大模型更快更省内存。

Gradient Multi-Normalization for Stateless and Scalable LLM Training

  • 通过多范数梯度归一化设计无状态优化器,避免存储历史信息。
  • 在10亿参数模型上比Adam快3倍,内存消耗显著降低。
  • 原理清晰可解释,适合大规模语言模型训练场景。

训练大语言模型通常依赖Adam等自适应优化器,这些优化器需存储额外状态以加速收敛,但带来巨大内存开销。近期工作SWAN(Ma et al., 2024)通过多步梯度预处理实现无状态优化,在性能上接近Adam。受此启发,本文提出一种新的无状态优化框架,基于多重范数对随机梯度进行归一化,并设计简单交替算法实现该归一化。理论上证明该方法可逼近任意精度的固定点,且SWAN是其特例,揭示了其设计本质。然而SWAN的白化/正交化步骤计算成本高,限制实际应用。本文基于这一理论视角,提出更高效、可扩展的无状态优化器:放松部分约束,大幅降低计算开销,同时保持内存效率。实验表明,在高达10亿参数的LLaMA模型预训练中,该方法相较Adam提速3倍,显著优于其他内存高效基线。

原文摘要 · Abstract (English)

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.

大模型训练无状态优化内存效率梯度归一化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。