跳转至

文章背景与核心概要

在训练混合专家（MoE）模型时，优化器状态内存往往是整个内存预算中占比最高的部分，这严重制约了模型规模的扩展。例如，对于一个拥有 6.78B 参数、占用 12.6 GB bfloat16 权重的 MoE 语言模型，标准 AdamW 仅仅为了跟踪一阶和二阶矩就需要消耗高达 50.6 GB 的内存。

为了解决这一痛点，本文作者 Nuemaan Malik 提出了 SkewAdam 优化器。该方法的核心洞察在于：MoE 的三个核心参数群体——密集主干（dense backbone）、专家（experts）和路由（router）——在规模和梯度统计特征上存在巨大差异，因此不应分配相同的优化器状态。通过采用分层状态分配（tiered state allocation）策略，SkewAdam 在将优化器状态内存大幅缩减超过 97% 的同时，其验证困惑度（validation perplexity）表现甚至超越了标准的基准优化器。

优化器状态究竟该存放在哪里？面向内存高效的混合专家模型训练的分层状态分配

作者： Nuemaan Malik
提交时间： 2026年7月21日；最后修订： 2026年8月13日（v2版）
学科领域： 机器学习（cs.LG）；人工智能（cs.AI）
arXiv 链接： 2607.19058 [cs.LG] | DOI： 10.48550/arXiv.2607.19058
资源： 代码与训练日志

📌 执行摘要

Training Mixture-of-Experts (MoE) models is severely bottlenecked by optimizer state memory, which often constitutes the largest single item in the memory budget. For instance, standard AdamW requires 50.6 GB of memory just to track first and second moments for a 6.78B-parameter MoE language model with 12.6 GB of bfloat16 weights.

训练混合专家（MoE）模型严重受制于优化器状态内存的瓶颈，这通常占到整个内存预算中最大的一项。例如，标准 AdamW 仅为了跟踪一个 6.78B 参数、具有 12.6 GB bfloat16 权重的 MoE 语言模型的一阶和二阶矩，就需要 50.6 GB 的内存。

This paper introduces SkewAdam, a novel optimizer designed around the insight that the three core parameter populations of an MoE—the dense backbone, the experts, and the router—differ vastly in size and gradient statistics, and therefore should not be allocated identical optimizer states. By utilizing a tiered state allocation strategy, SkewAdam reduces optimizer state memory by over 97% while simultaneously outperforming standard baselines in validation perplexity.

本文引入了一种名为 SkewAdam 的新型优化器。该优化器的设计基于这样一个深刻见解：MoE 的三个核心参数群体——密集主干（dense backbone）、专家（experts）和路由（router）——在规模和梯度统计上差异巨大，因此不应分配相同的优化器状态。通过采用分层状态分配（tiered state allocation）策略，SkewAdam 将优化器状态内存减少了 97% 以上，同时在验证困惑度上超越了标准基准。

🏗️ 分层分配策略（SkewAdam）

SkewAdam divides model parameters into three distinct populations and tailors the optimizer state allocation to each:

SkewAdam 将模型参数划分为三个不同的群体，并为每个群体量身定制了优化器状态分配方案：

  1. Dense Backbone (5% of parameters): Allocated float32 momentum plus a factored second moment.
  2. Experts (95% of parameters): Allocated a factored second moment alone (no momentum).
  3. Router (<0.01% of parameters): Allocated an exact second moment for precise routing control.
  1. 密集主干（占参数的 5%）： 分配 float32 动量加上因式分解的二阶矩。
  2. 专家（占参数的 95%）： 仅分配因式分解的二阶矩（无动量）。
  3. 路由（占参数的 <0.01%）： 分配精确的二阶矩以实现精确的路由控制。

内存影响

  • Standard AdamW State: 50.6 GB
  • SkewAdam State: 1.29 GB (only 2.6% of AdamW)
  • Peak Training Memory: Drops from 81.4 GB to 31.3 GB, bringing large-scale MoE training comfortably within the budget of a single 40 GB accelerator.
  • 标准 AdamW 状态： 50.6 GB
  • SkewAdam 状态： 1.29 GB（仅为 AdamW 的 2.6%）
  • 训练峰值内存： 从 81.4 GB 下降至 31.3 GB，使得大规模 MoE 训练可以轻松在单张 40 GB 加速卡的预算范围内进行。

📊 实证结果与发现

In controlled evaluations across 82M tokens from identical initializations, SkewAdam achieves superior performance compared to popular baseline optimizers:

在从相同初始化开始、基于 82M Token 的控制评估中，SkewAdam 相比流行的基准优化器取得了更优异的性能：

  • SkewAdam: 108.4 validation perplexity (and settles router load balance to within 1% of its uniform floor)
  • AdamW: 126.8 validation perplexity
  • Muon: 120.2 validation perplexity
  • Lion: 393.7 validation perplexity
  • SkewAdam： 108.4 验证困惑度 （并且将路由负载均衡稳定在其均匀下限的 1% 以内）
  • AdamW： 126.8 验证困惑度
  • Muon： 120.2 验证困惑度
  • Lion： 393.7 验证困惑度

消融实验的关键要点

  • Memory vs. Accuracy: A tier ablation study revealed that carrying twenty times the state yields the same perplexity, proving that the tiers buy memory efficiency rather than accuracy.
  • What Drives Accuracy? Platform runs isolating specific components revealed that:
  • Removing momentum entirely costs 31 perplexity points (tuned Adafactor: 139.7).
  • Replacing the factored second moment and its update clipping with a full second moment costs 10 perplexity points (tuned AdamW: 118.5).
  • Consequently, neither tuned baseline matches the untuned tiered policy, demonstrating that where optimizer state lives is just as critical as how much of it exists.
  • 内存与准确率的关系： 分层消融研究表明，携带二十倍的状态会产生相同的困惑度，这证明这些层带来的收益是内存效率而非准确率的提升。
  • 是什么驱动了准确率？ 隔离特定组件的平台运行揭示了：
  • 完全移除动量会导致 31 个困惑度点的损失（经过调优的 Adafactor：139.7）。
  • 用完整的二阶矩替换因式分解的二阶矩及其更新裁剪会导致 10 个困惑度点的损失（经过调优的 AdamW：118.5）。
  • 因此，无论是哪种经过调优的基准都无法匹敌未经调优的分层策略，这表明优化器状态存放在哪里与它存在多少同样至关重要。

🔗 快速链接与全文访问