提出异步LMO优化方法,提升分布式训练效率。
Ringmaster LMO: Asynchronous Linear Minimization Oracle Momentum Method
- 基于延迟阈值机制,允许异步更新
- 在异构系统下实现最优时间复杂度
- 适合大规模异步分布式训练场景
Muon 近来成为训练神经网络的有力替代方案,其在大规模预训练中表现优异,且矩阵结构更新在实践中速度更快。然而,现有基于线性最小化算子(LMO)的方法通常采用同步方式,在异构分布式系统中,因各计算节点完成梯度计算的时间不一,同步训练需频繁等待慢速节点,导致效率下降。本文提出 Ringmaster LMO,一种用于无约束随机非凸优化的异步 LMO 动量方法。该方法借鉴 Ringmaster ASGD 的延迟阈值思想,通过丢弃过时梯度实现最优时间复杂度。理论方面,在广义 (L₀, L₁) 平滑条件下建立收敛性保证,并设计无需调参、步长递减且自适应延迟阈值的变体。进一步将迭代保证转化为异构计算时间下的时间复杂度界。在经典欧氏光滑设置下,该界恢复了 Ringmaster ASGD 的最优时间复杂度。实验在随机二次问题和 NanoChat 语言模型预训练上验证:随着系统异构性增加,Ringmaster LMO 的优势愈发明显,显著优于强同步与异步基线方法。
原文摘要 · Abstract (English)
Muon has recently emerged as a strong alternative to AdamW for training neural networks, with encouraging large-scale pretraining results and growing evidence that matrix-structured updates can be faster in practice. Yet Muon, and more generally Linear Minimization Oracle (LMO) based methods, are typically used synchronously. This is problematic in heterogeneous distributed systems, where workers complete gradient computations at different speeds and synchronous training must repeatedly wait for slower workers. In this work, we introduce Ringmaster LMO, an asynchronous LMO-based momentum method for unconstrained stochastic nonconvex optimization. Our method builds on the delay-thresholding idea of Ringmaster ASGD. For SGD-type methods, Ringmaster ASGD achieves optimal time complexity by discarding overly stale gradients. Ringmaster LMO extends this mechanism to general LMO-based updates. We establish convergence guarantees under generalized $(L_0, L_1)$-smoothness and further develop a parameter-agnostic variant with decreasing stepsizes and adaptive delay thresholds. Finally, we translate our iteration guarantees into time complexity bounds under heterogeneous worker computation times. In the classical Euclidean smooth setting, these bounds recover the optimal time complexity of Ringmaster ASGD. Experiments on stochastic quadratic problems and NanoChat language-model pretraining show that the advantages of Ringmaster LMO grow with system heterogeneity and that the method outperforms strong synchronous and asynchronous baselines.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。