StoSignSGD让符号梯度下降稳定训练大模型,突破精度极限。
StoSignSGD: Unbiased Structural Stochasticity Fixes SignSGD for Training Large Language Models
- 在符号梯度中引入结构化随机性,保持更新无偏。
- FP4下仍稳定训练,比基线快2.14倍,准确率提升1.13点。
- 适合超低精度预训练与70亿参数模型微调,可解释性强。
基于符号的优化算法(如SignSGD)因在分布式学习和大模型训练中的优异表现而受到关注。尽管如此,其在非光滑目标上易发散,而这类目标广泛存在于含ReLU、最大池化及专家混合层的模型中。为此,我们提出StoSignSGD,通过在符号算子中注入结构化随机性,在保持无偏更新的同时解决收敛问题。在凸优化框架下,StoSignSGD严格解决了SignSGD的不收敛问题,达到匹配理论下界的收敛速率。对于更复杂的非凸非光滑优化,我们引入广义驻点度量,证明了该算法在复杂度界上优于现有最优结果,且优势随维度增长。实验表明,StoSignSGD在多种大语言模型(LLM)训练场景中稳定高效:在极端低精度预训练(包括FP8与更具挑战性的FP4)下,当AdamW完全崩溃时,该算法依然稳定并持续领先;在FP8下速度提升1.44至2.14倍;在4比特精度下,对最大规模的OLMo2-370M模型,下游任务准确率超越最强稳定基线1.13点,且优势随模型与数据规模扩大而增强。在70亿参数模型的数学推理微调任务中,也显著优于AdamW与SignSGD。为揭示其有效性,我们构建了通用优化器转符号形式的转换框架,通过系统消融验证核心设计组件。
原文摘要 · Abstract (English)
Sign-based optimization algorithms, such as SignSGD, have garnered attention for their performance in distributed learning and training large foundation models. Despite their empirical superiority, SignSGD is known to diverge on non-smooth objectives, which are ubiquitous due to ReLUs, max-pools, and mixture-of-experts. To overcome this limitation, we propose StoSignSGD, an algorithm that injects structural stochasticity into the sign operator while maintaining an unbiased update step. In the regime of (online) convex optimization, StoSignSGD rigorously resolves the non-convergence issues of SignSGD, achieving a sharp convergence rate matching the lower bound. For the more challenging non-convex non-smooth optimization, we introduce generalized stationary measures that encompass prior definitions, proving that StoSignSGD improves upon the best-known complexity bounds by dimensional factors. Empirically, StoSignSGD is stable and efficient across diverse large language model (LLM) training regimes. In aggressive low-precision pretraining, which spans both FP8 and the far more demanding FP4 regime where AdamW fails catastrophically, StoSignSGD stays stable and consistently performs the best. It attains a 1.44x to 2.14x speedup over established baselines under FP8. Under 4-bit precision, it improves downstream accuracy on the largest OLMo2-370M model by 1.13 points over the strongest stable baseline, and this advantage grows as both the model size and the data scale up. When fine-tuning 7B LLMs on mathematical reasoning tasks, StoSignSGD also delivers clear gains over both AdamW and SignSGD. Finally, to explain why it works, we develop a sign conversion framework that turns any general optimizer into its unbiased, sign-based counterpart. Using this framework, we decompose the core components of StoSignSGD and run a comprehensive ablation study to validate our design choices.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。