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

TL;DR SkewAdam 根据 MoE 中稠密骨干、专家和路由器在参数量与梯度分布上的差异，为不同参数群体分配不同精度的优化器状态，将显存降至 AdamW 的 2.6%，且困惑度超越多个主流优化器，揭示了优化器状态“存放位置”与精度同样关键。

问题

问题背景

稀疏混合专家 (MoE) 架构通过为每个 token 仅激活少量专家，实现了参数量与计算量的解耦，成为扩展大模型的主流范式。然而，训练过程中的优化器状态（如 Adam 的一阶动量与二阶方差）并未受益于稀疏激活——其对所有参数均需维护完整的状态副本，成为内存预算的最大单一开销。

现有方法的局限

主流优化器 AdamW 为每个参数保存 float32 格式的动量与二阶矩，在 MoE 场景下状态内存可达权重的数倍，严重限制可训练模型规模。内存高效方案如 Adafactor 使用因式分解的估计器削减二阶矩状态，但其完全丢弃一阶动量，导致收敛速度与最终精度明显落后。低精度优化器 (如 8-bit Adam) 虽能压缩状态，却会引入量化噪声，在 MoE 的梯度异质性下可能加剧不稳定。符号类方法 (如 Lion, Muon) 仅用梯度符号更新，内存极省但损失了二阶自适应信息，在稀疏路由的高方差梯度场景下难以稳定收敛。这些方法均采用统一状态分配策略——对所有参数一视同仁，忽略了 MoE 中稠密主干、专家与路由三类参数在规模与梯度统计上的本质差异。

技术挑战与重要性

MoE 的参数群体呈现极度偏斜分布：约 95% 的参数隶属专家层，但其梯度稀疏且噪声大；稠密骨干仅占 5%，却对全局表征至关重要；路由参数极少但需精准学习。为不同群体分配相同的优化器状态，实质上是在内存效率与收敛质量之间做不必要妥协。在单卡 40–80 GB 的硬件约束下，若无法缩减优化器内存，每 token 可训练的参数容量将触及天花板，直接阻碍 MoE 模型向更大规模、更强性能的演进。因此，设计状态非对称分配策略，在显著降低内存的同时保持甚至提升收敛效果，成为工程落地与前沿探索的共同焦点。

行业类比

如同在大规模推荐系统中，为高频特征保留完整统计信息、而对长尾特征采用压缩近似，这种按需分配状态的思路可在内存与精度间求得最优权衡。

核心洞察

• **分层状态分配**：MoE 模型的参数群体（稠密骨干、专家、路由器）在规模与梯度统计上差异显著，统一分配优化器状态会导致巨大内存浪费。SkewAdam 针对三者特点分别配置状态：骨干用 float32 动量+因子化二阶矩，专家仅用因子化二阶矩，路由器保留精确二阶矩。这一设计使状态内存从 AdamW 的 50.6 GB 降至 1.29 GB（2.6%），单卡 40 GB 即可训练，且验证困惑度反而优于全状态基线。与 Adafactor 对所有参数无差别使用因子化二阶矩、或低精度优化器对整个状态统一量化不同，SkewAdam 的分层策略考虑了参数角色的结构性差异，在内存与精度之间达到更优权衡，证明了优化器状态分配位置的重要性不低于状态总量。
• **动量保留是性能核心**：分层状态分配节省了内存，但消融实验表明，仅靠节省内存的分配方案并未带来困惑度收益——使用 20 倍状态的等分配置可匹配 SkewAdam 的性能。真正的性能提升来自为骨干参数保留动量（第一矩）。相比之下，Adafactor 虽然也使用因子化二阶矩，但完全丢弃动量，导致困惑度比 AdamW 落后 40 点以上。这一发现挑战了“因子化二阶矩足以替代完整 Adam”的流行假设，强调动量在适应损失景观曲率与梯度噪声中的关键作用。对工程实践的启示是：在 MoE 训练内存优化中，优先保留骨干参数的动量比全局统一压缩状态更为有效，为后续优化器设计指明了取舍方向。

方法

输入：MoE 参数群体的划分

SkewAdam 将 Mixture-of-Experts 模型的参数分为三个异构群体：

• dense backbone（约占 5% 参数），负责共享表示，梯度统计特性与密集模型相似。
• experts（约占 95% 参数），稀疏激活，各专家更新频次和梯度分布差异大。
• router（<0.01% 参数），决定 token 路由，对稳定性要求极高。

这三个群体在规模、梯度幅值及二阶矩估计需求上差异显著，为差异化状态分配提供了依据。

关键模块：分层状态分配与预估器设计

SkewAdam 的核心是为每个参数群体定制优化器状态，实现内存精确分配：

• 骨干网络：保留 float32 一阶动量（类似 Adam 的指数移动平均），并采用 factored 二阶矩估计（将二阶矩矩阵分解为行向量与列向量外积），既保留动量加速收敛，又通过分解降低内存占用。
• 专家网络：仅使用 factored 二阶矩估计，完全丢弃一阶动量。由于专家参数体量巨大且稀疏激活，动量的边际收益有限，而分解二阶矩可节省大量内存。
• 路由器：使用 exact 二阶矩（完整存储），不进行分解。尽管二阶矩精确计算内存需求稍高，但路由器参数极少，状态开销可忽略，保障路由负载均衡的稳定性。

此外，优化器统一采用 更新裁剪（update clipping）和 低精度更新（bfloat16 参数 + float32 状态），进一步控制训练数值范围与显存带宽。

输出：内存高效且收敛无损的优化器

最终优化器状态总大小仅 1.29 GB（AdamW 的 2.6%），峰值训练内存从 81.4 GB 降至 31.3 GB，可放入 40 GB 加速器。在 82M token 验证中，SkewAdam 达到 108.4 的 validation perplexity，优于 AdamW（126.8）和 Adafactor（139.7），并保持路由器负载均衡接近均匀分布。

与同类方法的差异：SkewAdam 首次将优化器状态的位置（即分配给哪类参数）与预估器类型解耦，证明分层分配比均匀压缩（如 AdamW）或全局简化（如 Adafactor）更能兼顾内存与收敛。

实验

实验设计

在一个 6.78B 参数 的 MoE 语言模型上，从 相同初始化 出发训练 82M tokens，对比四种优化器：SkewAdam、AdamW、Muon、Lion。评估验证困惑度、优化器状态内存、峰值训练内存及路由负载均衡。此外，通过 tier 消融实验 分析状态分配对精度的影响，并对基线 AdamW 和 Adafactor 进行学习率调优以测试上限。

关键发现

• 极致内存节省：SkewAdam 的优化器状态仅占 1.29 GB，为 AdamW（50.6 GB）的 2.6%；峰值训练内存从 81.4 GB 降至 31.3 GB，使得 40 GB 加速器 即可完成训练。
• 精度领先：在验证困惑度上达到 108.4，优于 AdamW（126.8）、Muon（120.2）和 Lion（393.7），路由负载均衡保持在均匀基线的 1% 以内。
• 动量是关键：消融实验表明，节省内存的 tier 分配 本身不会损害精度（同样配置用 20 倍状态可匹配），精度增益来自保留 动量（momentum）；而 Adafactor 丢弃动量后困惑度高 40 点。
• 调优无法弥补差距：对基线进行最优学习率搜索后，AdamW 最佳达 118.5，Adafactor 达 139.7，仍未超越 SkewAdam。

对比解读

SkewAdam 挑战了“优化器状态总量决定内存”的朴素认知，证明 状态如何分配与总量同等重要。在 MoE 中，dense backbone（~5% 参数）、expert（~95%）和 router（<0.01%）的梯度统计特性差异显著。通过为 backbone 保留 float32 动量 + 因子化二阶矩，为 expert 仅保留 因子化二阶矩，为 router 保留 精确二阶矩，实现了 针对性压缩，避免了精度损失。这对工程实践的启示是：不应一刀切地使用低精度优化器，而应根据参数角色的重要性和梯度特性，设计分层状态策略，从而在有限硬件预算下扩展模型规模。

行业影响

落地场景

SkewAdam 的分层优化器状态策略直接利于大规模 MoE 模型的训练与微调，尤其适合显存受限的环境：

• 云 GPU 服务：用更廉价的 40 GB 级加速器（如 A40、RTX 6000 Ada）替代 80 GB 设备训练千亿参数 MoE，降低算力门槛。
• 边缘与私有化部署：在医疗影像、金融风控等数据敏感场景，客户可在自有小 GPU 集群上微调 MoE 专家模型，无需上传数据。
• 连续学习与在线更新：对推荐系统、自动驾驶等需持续增量训练的业务，SkewAdam 大幅降低内存压力，使在线专家更新成为可能。

商业价值

成本削减是核心商业回报：

• 训练成本直降 60%+：显存占用从 81.4 GB 降至 31.3 GB，单卡即可训练原来需要更高规格或更多 GPU 并行的模型，硬件投入与能耗同步降低。
• 迭代速度提升：在固定预算下可并行更多实验，加速模型选优与上线周期。
• 产品形态扩展：得以在端侧或边缘设备上运行原本需要云端大卡的 MoE 推理或微调服务，为端智能产品提供新的技术可行性。

与现有产品 / 工作流的接口

SkewAdam 作为优化器，与现有训练栈无缝集成：

• 直接替换 PyTorch 中的 AdamW 或 Adafactor，仅需在初始化时对 dense/expert/router 参数分组并指定对应的状态策略。
• 兼容主流分布式训练框架（如 DeepSpeed、FSDP），配合激活检查点（activation checkpointing）可进一步压缩内存。
• 可封装进 Hugging Face Trainer 或自定义训练脚本，无需改变数据加载、模型架构，迁移成本极低。

具体落地 Use Case

1. 跨国电商推荐系统
   基于 MoE 的千亿参数推荐模型，需每日根据用户行为增量训练。使用 SkewAdam，在现有 40 GB GPU 集群上即可维持日更节奏，避免扩容到 80 GB GPU，年节省百万级硬件开支。

2. 自动驾驶多任务感知
   车载域控算力有限，当需要针对新场景微调感知专家时，SkewAdam 使车载计算平台（如 NVIDIA Orin）能以极低显存开销完成模型快速适配，无需将数据上传云端，满足实时性与隐私要求。

局限

• 实验仅在单个 6.78B 参数的 MoE 语言模型上进行，且仅训练了 82M tokens；未在更大规模（如数十亿参数、数万亿 tokens）或不同架构的 MoE 上验证，因此层级分配策略的泛化性和扩展性尚不明确。此外，模型架构可能对结果有较大影响，缺乏不同 MoE 变体（如 Expert Choice Routing 等）的测试。
• 与基线方法的对比不够全面：主要比较了 AdamW、Adafactor、Muon 和 Lion，但缺少与近期为 MoE 或大模型专门设计的低内存优化器（例如 Galore 已被提及但不够深入，以及 8‑bit Adam 等低精度方案）的详细对比，也未探讨与如 ZeRO‑Infinity 等系统级内存优化相结合的潜力。
• 方法需要根据参数组的人工经验划分层级（dense backbone、experts、router），并为每个层级选择具体的状态配置（如是否保留 momentum、使用 factored second moment 等），这种设计难以自动适配其他模型结构或训练任务，超参数搜索空间大，可能成为实际部署中的工程负担。