MoE 训练显存吃紧?MESH 方法砍掉六成优化器内存,损失只差一点,做大规模训练的可以试试。
论文提出 MESH,一种用于 MoE 专家层的隐藏动量 Sinkhorn 更新。在 110M 参数 DeepSeek 风格 MoE 预训练中,SAGE/Sinkhorn 混合把优化器状态从 0.883GB 降到 0.331GB,但评估损失 3.8265,差于 AdamW 基线的 3.58-3.64。MESH 在梯度缓冲区生命周期内恢复一阶矩信号,且不将专家一阶矩存入优化器状态。两个额外种子中,MESH 和 MESH-B 相对 AdamW 减少 62.5% 优化器状态内存,峰值 CUDA 分配下降约 12.6%,评估损失差距较小。
MESH: Memory-Efficient Sinkhorn Optimization for Mixture-of-Experts Training
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.