Where Should Optimizer State Live? Tiered State Allocation for Memory-Efficient Mixture-of-Experts Training
SkewAdam 按 MoE 参数角色分层配置优化器状态:骨干保留 FP32 动量和分解二阶矩,专家仅保留分解二阶矩,路由器使用精确二阶矩。6.78B 模型的状态从 AdamW 的 50.6GB 降至 1.29GB,峰值训练显存由 81.4GB 降至 31.3GB;训练 8200 万 token 后验证困惑度为 108.4,优于 AdamW 的 126.8。
INSIGHTMoE 训练的优化器状态不应按参数一刀切:占参数 95% 的专家层最适合压缩状态,而不足 0.01% 的路由器值得保留精确统计。对需要在 40GB 级加速卡上训练领域 MoE 的团队,这种按角色分配精度的方式比统一量化更易控制收敛风险。