提出一种内存高效优化方法,收敛速度比现有方法快
Omni-Masked Gradient Descent: Memory-Efficient Optimization via Mask Traversal with Improved Convergence
- 通过掩码遍历策略减少内存占用,实现高效训练
- 在非凸环境下达到约 ε⁻³ 的迭代复杂度,优于传统 O(ε⁻⁴)
- 可无缝集成主流优化器,适合大模型训练场景
为缓解大语言模型训练中的显存瓶颈,内存高效优化方法受到广泛关注。现有方法或缺乏明确的收敛保证,或在非凸设置下仅能达到标准的 O(ε⁻⁴) 迭代复杂度。本文提出基于掩码遍历的内存高效优化方法 Omni-Masked Gradient Descent (OMGD),并提供了非凸收敛分析,证明其可达到严格改进的迭代复杂度 Õ(ε⁻³),用于寻找 ε-近似驻点。实验表明,OMGD 是一种轻量级、即插即用的方法,能无缝集成到多数主流优化器中,在微调和预训练任务中均持续优于对比基线。
原文摘要 · Abstract (English)
Memory-efficient optimization methods have recently gained increasing attention for scaling full-parameter training of large language models under the GPU-memory bottleneck. Existing approaches either lack clear convergence guarantees, or only achieve the standard ${\mathcal{O}}(ε^{-4})$ iteration complexity in the nonconvex settings. We propose Omni-Masked Gradient Descent (OMGD), an optimization method based on mask traversal for memory efficient training, and provide a nonconvex convergence analysis that establishes a strictly improved iteration complexity of $\tilde{\mathcal{O}}(ε^{-3})$ for finding an $ε$-approximate stationary point. Empirically, OMGD is a lightweight, plug-and-play approach that integrates seamlessly into most mainstream optimizers, yielding consistent improvements over competitive baselines in both fine-tuning and pre-training tasks.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。