论文页面 - 优化器状态该驻留何处？面向内存高效的混合专家模型训练的分层状态分配

来源：https://huggingface.co/papers/2607.19058

摘要

优化器状态是混合专家模型（MoE）训练中内存预算最大的单项开支：对于一个67.8亿参数的MoE语言模型，AdamW需要占用50.6 GB的存储空间来保存一阶和二阶矩，以更新12.6 GB的bfloat16权重。我们研究了SkewAdam优化器，它基于以下观察构建：MoE的三个参数群体——密集主干网络、专家网络和路由网络——在规模和梯度统计特性上差异足够大，不应共享相同的状态。SkewAdam为主干网络（占参数量的5%）保留float32动量加上因式分解的二阶矩，为专家网络（占95%）仅保留因式分解的二阶矩，而为路由网络（<0.01%）保留精确的二阶矩。由此产生的状态占用1.29 GB，仅为AdamW的2.6%，峰值训练内存从81.4 GB降至31.3 GB，落入了40 GB加速器的预算范围内。在从相同初始条件出发、基于8200万tokens的受控对比实验中，SkewAdam的验证困惑度达到108.4，领先于AdamW（126.8）、Muon（120.2）和Lion（393.7），并将路由负载均衡度稳定在其均匀下限的1%以内。这种分配策略并非取得该困惑度的关键：对分层方案进行消融实验，在状态量扩大二十倍时仍能达到相同效果；而同样采用因式分解估计量但舍弃动量的Adafactor，则落后40点。分层方案在没有牺牲准确性的前提下节省了内存；真正的准确性来源于保留动量，而这在统一优化器中同样适用。调整基线学习率可以缩小差距，但无法完全弥合：经过最佳调参的AdamW达到118.5，调参后的Adafactor达到139.7。这些结果表明，优化器状态存放在何处，与状态总量的多少至少同等重要。

查看arXiv页面 (https://arxiv.org/abs/2607.19058)查看PDF (https://arxiv.org/pdf/2607.19058)GitHub1 (https://github.com/nuemaan/skewadam)添加到收藏 (https://huggingface.co/login?next=%2Fpapers%2F2607.19058)

在您的Agent中获取本文：

hf papers read 2607.19058

没有最新的CLI？curl -LsSf https://hf.co/cli/install.sh | bash

引用本文的模型 0

尚无模型关联本文

请在模型的README.md中引用 arxiv.org/abs/2607.19058，以便在此页面建立链接。

引用本文的数据集 0

尚无数据集关联本文

请在数据集的README.md中引用 arxiv.org/abs/2607.19058，以便在此页面建立链接。

引用本文的Space 0

尚无Space关联本文

请在Space的README.md中引用 arxiv.org/abs/2607.19058，以便在此页面建立链接。

包含本文的收藏集 0

尚无收藏集包含本文

将本文添加到一个收藏集 (https://huggingface.co/new-collection) 以便在此页面建立链接。