提出MESH优化器,让MoE模型训练内存减少62.5%且损失仅小幅上升。
MESH: Memory-Efficient Sinkhorn Optimization for Mixture-of-Experts Training
- 用隐藏动量机制恢复专家梯度的时间平滑信号,不存完整状态
- 在110M参数模型上将优化器内存从0.883GB降至0.331GB
- 适合大模型训练中追求内存效率的工程师和研究者
内存高效的矩阵优化器如Sinkhorn梯度下降能移除密集Transformer矩阵的AdamW状态,但直接用于混合专家(MoE)训练时不可靠。我们在一个110M参数的nanowhale DeepSeek风格MoE预训练设置中研究该问题。采用SAGE/Sinkhorn混合方案可将优化器状态从0.883GB降至0.331GB,但评估损失升至3.8265,远高于同环境下AdamW基线(3.58–3.64)。我们发现路由后的MoE专家矩阵是主要失败点:其梯度具有条件性、时间变化性,且难以被无状态的Sinkhorn归一化满足。为此提出MESH——一种针对MoE专家的隐式动量Sinkhorn更新机制。MESH通过梯度缓冲生命周期恢复时间一阶矩信号,无需存储专家一阶矩状态。该方法为可选的块预条件变体,引入粗粒度神经元/块逆均方根缩放因子。消融实验表明,矩阵归一化前的时间平滑是核心因果因素;块/神经元预条件可进一步提升内存-性能权衡,但非普遍必需。在两个额外种子下,MESH与MESH-B相较AdamW分别降低62.5%的优化器内存和约12.6%的PyTorch CUDA峰值分配,评估损失仅小幅增加。全状态诊断变体在消融中恢复了接近AdamW的性能,支持结论:MoE专家需要时间平滑,但无需完整的坐标级AdamW状态。
原文摘要 · Abstract (English)
Memory-efficient matrix optimizers such as Sinkhorn gradient descent remove most AdamW optimizer state for dense Transformer matrices, but direct application to Mixture-of-Experts (MoE) training is unreliable. We study this failure in a controlled 110M-parameter nanowhale DeepSeek-style MoE pretraining setting. A SAGE/Sinkhorn hybrid reduces optimizer state from 0.883GB to 0.331GB but degrades evaluation loss to 3.8265, far above the AdamW baselines observed in the same setup (3.58--3.64 across the seeds we study). We show that routed MoE expert matrices are the dominant failure point: their gradients are conditional, temporally varying, and poorly served by stateless Sinkhorn normalization. We propose MESH, a hidden-momentum Sinkhorn update for MoE experts. MESH restores a temporal first-moment signal through the gradient-buffer lifecycle, without storing the expert first moment as optimizer state. MESH is an optional block-preconditioned variant that adds a coarse neuron/block inverse-RMS multiplier. Across ablations, temporal smoothing before matrix normalization is the primary causal ingredient; block/neuron preconditioning can improve the memory-quality frontier, but is not established as universally necessary. In two additional seeds, MESH and MESH-B reduce optimizer-state memory by 62.5\% and peak PyTorch CUDA allocation by about 12.6\% relative to AdamW, with a modest evaluation-loss gap. Full-state diagnostic variants recover AdamW-like performance in ablations, supporting the conclusion that MoE experts need temporal smoothing, but not necessarily full coordinate-wise AdamW state.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。