← 返回 PaperDaily
大模型与智能体
仅用AdamW 2.6%状态！SkewAdam分层分配让67亿参数MoE单卡可训
还在为训练大模型时显存爆掉而烦恼？AdamW为671亿参数的MoE模型准备了50.6GB的优化器状态，比模型权重本身还大4倍！独立研究者Nuemaan Malik发现，MoE模型的不同部分（骨干网络、专家层、路由器）对优化器状态的需求截然不同。SkewAdam通过分层分配，将优化器状态压缩到1.29GB（仅占AdamW的2.6%），峰值内存降至31.3GB，
龙哥读论文
阅读 3
查看原文
龙哥推荐理由：
还在为训练大模型时显存爆掉而烦恼？AdamW为671亿参数的MoE模型准备了50.6GB的优化器状态，比模型权重本身还大4倍！独立研究者Nuemaan Malik发现，MoE模型的不同部分（骨干网络、专家层、路由器）对优化器状态的需求截然不同。SkewAdam通过分层分配，将优化器状态压缩到1.29GB（仅占AdamW的2.6%），峰值内存降至31.3GB，单张40GB GPU就能训6.78B参数模型，且验证困惑度反而更低（108.4 vs 126.8）。这篇论文的价值不仅在于一个具体的优化器，更在于一个设计原则：让优化器状态各司其职，而非一视同仁。
原论文信息如下：
优化器状态也有“偏科”？MoE训练的内存困境
训练大模型，显存是永远的痛。尤其是混合专家模型（MoE），虽然参数量大得吓人，但每个token只激活一小部分专家——比如一个6.78B参数的MoE，每token只激活约440M参数，计算量对应小模型。然而，优化器状态却“一视同仁”：AdamW为每个参数保存两个float32动量，总共8字节，而参数本身只有2字节（bfloat16）。所以，6.78B参数的MoE，权重只占12.6GB，优化器状态却要吃掉50.6GB！整个训练峰值内存高达81.4GB，直接告别单卡40GB的散热。
更苦的是，这50.6GB的状态主要是为专家层准备的——专家层占了95%的参数（6.4B）。但每个专家在top-2路由下，平均只处理约1/64的token。也就是说，你拿着float32的动量去平滑一个几乎没怎么见过的梯度，纯属浪费。路由器（gate）只有0.52M参数，虽然小，但它的梯度携带负载均衡信号，若像Lion那样直接丢掉梯度幅度信号，路由会乱套。
现有的内存高效优化器，比如Lion、Muon、Adafactor，都把网络当成一个均匀的块来处理。Lion只保留一个动量缓冲，用符号更新，在路由器那里丢掉了梯度幅度；Muon对专家矩阵进行牛顿-舒尔茨正交化，但计算量很大；Adafactor虽然分解了二阶矩，但要么在所有地方都丢弃动量，要么在所有地方都保留，没有区分对待。它们都没问一个问题：模型的不同部分，真的需要同等水平的优化器状态吗？
SkewAdam：为MoE量身定制的分层分配策略
SkewAdam的核心洞察简单粗暴：MoE模型有三个参数群体，它们的大小和梯度统计差异巨大，不该用同一种优化器状态配方。论文把参数分为三层：骨干网络（backbone）、专家层（experts）、路由器（router）。下面这张图直观展示了分层分配方案：
骨干网络（~5%参数）包括嵌入层、注意力投影、稠密前馈块、LayerNorm。这些参数每个token都参与，梯度密集，所以float32动量在这里很值；同时保留一个分解的二阶矩（row/column统计）。这5%的参数占用了1.27GB的动量（加上二阶矩共约1.3GB，占优化器状态的主体）。
专家层（~95%参数）128个专家，每个专家矩阵4096x4096x3（SwiGLU有3个权重矩阵）。在top-2-of-128路由下，每个专家平均只处理约1/64的token，梯度稀疏且方差大。在这里保留float32动量需要24GB，但大部分梯度更新非常稀疏，动量平滑作用微弱。因此SkewAdam只保留分解的二阶矩（row和column统计），完全丢弃动量。Adafactor的前期工作表明，在有二阶矩自适应学习率的情况下，可以安全移除动量而不会明显损失效果。专家层的分解二阶矩仅需12MB（每个矩阵n+m个float32，共128个专家矩阵）。
路由器（<0.01%参数）128x4096的权重，仅0.52M参数。路由器决定了每个token去哪两个专家，其梯度的相对幅度携带负载均衡信号。如果像Adafactor那样分解二阶矩，会将不同logit的方差混在一起，破坏路由决策。所以SkewAdam保留完整的、非分解的float32二阶矩，仅占2MB。这2MB买来了每个logit的独立自适应学习率，是路由稳定性的关键。
注意，路由器不使用权重衰减（weight decay），因为权重衰减会干扰路由偏向。此外，所有参数都使用bfloat16主权重，更新通过带抖动随机舍入（dithered rounding）的float32计算后写入。这种低精度策略对所有优化器一视同仁，公平比较。
分解二阶矩的公式如下：对于矩阵W，梯度G，保持行统计R和列统计C，指数移动平均，然后重建二阶矩估计为 R*C^T / mean(R)。这种做法与Adafactor完全一致，但只用于专家层和骨干网络。
实验效果：用1/40的状态超越AdamW
论文在6.78B参数的MoE模型上进行对比实验，所有优化器从相同初始化开始，使用相同的数据顺序，训练82M token。下面这张表格和内存对比图直接展示差距：
表1：6.78B参数MoE，10,000步从共享初始化的结果。状态大小为理论计算（附录B），内存和吞吐量为最终步测量。负载均衡损失在均匀路由下理论上限为0.05。
从表中可以看到，SkewAdam以1.29GB的优化器状态（仅占AdamW的2.6%）实现了最好困惑度108.4，同时负载均衡损失0.0505，几乎达到了均匀分配的下限（0.05）。AdamW虽然也保持很好的负载均衡（0.0502），但内存是SkewAdam的39倍，困惑度却差了18.4。Lion彻底扑街，困惑度393.7，说明丢弃梯度幅度信号在MoE路由中是不可取的。Muon在后期负载均衡损失上升（0.0608），比SkewAdam高出22%，困惑度120.2也落后约12个点。
吞吐量方面，SkewAdam每秒处理5000个token，比AdamW快6.6%，比Lion略慢1.5%（Lion只有一个缓冲，更新开销最小），比Muon快47%。Muon因为要对6.4B的专家参数做牛顿-舒尔茨迭代，计算开销明显。
实战验证：非均匀分配带来内存与性能双赢
内存省了97.4%，那效果会不会崩？别急，论文设计了一整套精密的实验，从多个维度验证了SkewAdam不仅没崩，还反超了AdamW。
实验设置：论文训练一个6.78B参数的decoder-only transformer，有两个块：第一个块是dense SwiGLU前馈层，第二个块是MoE层（128个专家，宽度4096，top-2路由）。模型总参数量6.78B，但每token只激活约440M参数。所有优化器从相同的初始化出发，使用相同的数据顺序，训练82M token（10,000步）。数据来自OpenWebText，按95/5比例切分训练/验证集，验证集使用相同的64个批次。
所有优化器都使用bfloat16主权重和带抖动随机舍入（dithered rounding）的更新路径，保证公平。学习率使用cosine调度，3% warmup。AdamW和SkewAdam的学习率3e-4，Lion为1e-4（其作者推荐低3-10倍），Muon矩阵为0.02、嵌入/路由器/向量参数为内部Adam的1e-3。运行在单张NVIDIA H200 GPU上（141GB显存，只为了让AdamW跑起来）。
从表中可以清晰看到：SkewAdam以1.29GB的优化器状态（仅占AdamW的2.6%）实现了最好困惑度108.4，同时负载均衡损失0.0505，几乎达到了均匀分配的理论下限（0.05）。AdamW虽然也保持很好的负载均衡（0.0502），但内存是SkewAdam的39倍，困惑度却差了18.4。Lion彻底扑街，困惑度393.7，几乎是随机猜的水平。Muon在后期负载均衡损失上升（0.0608），困惑度120.2也落后约12个点。
吞吐量方面，SkewAdam每秒处理5000个token，比AdamW快6.6%，比Lion略慢1.5%（Lion只有一个缓冲，更新开销最小），比Muon快47%。Muon因为要对6.4B的专家参数做牛顿-舒尔茨迭代（Newton-Schulz iteration），计算开销明显。
负载均衡损失函数是Switch Transformer引入的辅助损失（Auxiliary loss），等于α * E * sum(p_i * f_i)，当路由完全均匀时等于α=0.05。从图3看，SkewAdam和AdamW从4000步后就在0.05附近（0.0505和0.0502），非常稳定。Lion始终偏高7-9%，说明负载均衡还行但模型没学语言。Muon在最后1000步突然跳高22%到0.0608，这可能是不稳定性的开始，因为它的困惑度也在那段时间进入平台期。
SkewAdam最需要回答的比较对象是Adafactor——它的分解二阶矩就是从Adafactor来的。论文在获得更多计算资源后，在H100 MIG切片上重新运行了比较：
SkewAdam在H100上复现为109.0（H200上是108.4，机器间差异仅0.5%），说明协议跨硬件可迁移。Adafactor状态仅12MB（不含动量），但困惑度卡在149.5，最后1000步提升不到0.1，比SkewAdam差了40个点！两者共享相同的rank-one二阶矩估计器和更新clipping，唯一的区别是SkewAdam保留了动量而Adafactor完全丢弃了动量（并使用了另一种衰减调度）。这说明动量在MoE训练中是不可或缺的——即便只在骨干网络上保留，也能带来巨大的困惑度改善。
GaLore-style（rank 128）则彻底失败（困惑度1839.9），路由负载均衡倒是很好（0.0510），但完全没学会语言模型。作者谨慎地表示这可能是自己实现的配置不当，但至少说明把稀疏专家梯度投影到低秩子空间需要非常小心。
论文在MI300X GPU上做了一组关键的消融实验，逐层切换策略：
结果非常惊人：四种配置的困惑度都在108.2-108.9之间（单种子噪声范围内），负载均衡都约0.0505。但优化器状态从1.29GB到25.29GB，差了二十倍！具体细节：
给专家加动量：状态从1.29GB飙升到25.29GB，峰值从31.4GB到55.4GB，困惑度只变了0.2（108.7 vs 108.9）——专家动量的确是完全的浪费。
分解路由器二阶矩：状态从1.29GB降到1.28GB，困惑度和负载均衡几乎不变（108.2 vs 108.9）——精确路由器二阶矩的2MB开销是无害的，但不做也不会明显变差。
均匀配置（所有层都加动量+分解二阶矩）：困惑度108.3，状态25.29GB——这就是"满配"但浪费的状态，性能与SkewAdam相同。
这个消融实验完美回答了"哪个层级贡献了性能"：贡献主要来自骨干网络上的动量，而专家层动量完全是死重。这也证明了分层策略的贡献是内存而非更好的优化质量——策略用1/20的状态达到了与满配相同的效果。
学习率扫描（Learning Rate Sweep）：打不过就是打不过
消融实验后，剩下一个可能的质疑：SkewAdam领先AdamW是否只是因为AdamW的学习率没调好？论文对两个最强基线做了学习率扫描：
调优确实帮AdamW从126.8降到了118.5，Adafactor从149.5降到了139.7。但SkewAdam（未调优）的108.4-109.0，领先最佳AdamW（118.5）约10个点（20倍种子级标准差），领先最佳Adafactor（139.7）约31个点。调优缩小了差距，但差距并没有消失。这再次印证：动量在MoE训练中的价值，以及分层分配在内存上的结构性优势。
有趣的是，AdamW和Adafactor的最佳学习率都是1e-4，而SkewAdam默认3e-4。AdamW-Aafactor之间约21个点的差距（118.5 vs 139.7）与表3消融实验的动量故事一致：Adafactor完全丢弃了动量。
零样本评估（Zero-shot Evaluation）
为了完整性，论文用LM Evaluation Harness在PIQA、WinoGrande、HellaSwag、ARC-Challenge上评估了最终checkpoint。
启示：设计原则比具体算法更重要
SkewAdam的贡献，论文自己定性得很清楚：不是发明了更好的优化器，而是提出了一种设计原则（Design Principle）—— 优化器状态应该在模型的哪些层级上、以什么精度存在，应该根据各层级的梯度统计特征来分配。
对于更深的多层MoE模型（如Mixtral 8x7B或DeepSeekMoE），分层分配策略可以直接应用，可能效果更显著（因为专家层更多，状态浪费更大）。
与其他内存压缩技术正交：SkewAdam的分层原则可以和8-bit优化器、ZeRO分片、梯度检查点等技术叠加，节省更多内存。
对于非MoE模型（如Dense模型）：虽然分层分配的最大受益者可能是MoE，但对Dense模型中的不同层级（如底层、中层、顶层，或者Embedding与其他层）也可以考虑差异化分配。
模型只有两个块，95%的参数集中在一个专家bank中，虽然适合压力测试，但多层MoE的情况未验证。大多数实验是单种子运行（主对比除外）。82M token和128上下文非常小，与生产环境差距巨大。权重衰减在所有运行中无效（因为学习率cosine+3e-4+80M token下权重衰减几乎无影响），扩展到更长训练时必须重新引入权重衰减并确保它在bfloat16更新中存活。零样本评估接近随机，不需要过度解读。
总的来说，这篇论文的价值不在于一个具体的优化器实现，而在于一个非常务实的洞察：你的模型不是均匀的，为什么你的优化器是？给每个部分分配它最需要的状态精度，可以大幅节省内存，而性能只升不降。
对于普通研究者，这个思路最直接的启发是：不要盲目套用AdamW的默认配置，想想你的模型哪些部分是梯度密集的、哪些是稀疏的，然后可对优化器状态做一些精细化分配。可能不需要做到SkewAdam这么极端，但骨干网络保留动量、专家层用AdaFactor、路由器用精确二阶矩这种"混搭"思路，绝对值得一试。
最后，代码已开源（MIT协议），单文件实现，无额外依赖，可以直接集成到你的训练脚本中体验一下。
龙迷三问
SkewAdam 和 Adafactor 到底有什么区别？SkewAdam用在哪里？SkewAdam 和 Adafactor 共享相同的分解二阶矩估计方法（rank-one factorization），但有两个关键区别。第一，SkewAdam 在骨干网络（backbone）保留了完整的 float32 动量，而 Adafactor 在论文的对比中完全丢弃了动量（并使用了不同的衰减调度）。第二，SkewAdam 根据参数角色（骨干网络、专家、路由器）分配不同的状态组合，而 Adafactor 对所有参数一视同仁（要么全动量要么全无）。在 MoE 的训练中，骨干网络的动量非常重要（实验证明骨干网络的动量贡献了与满配状态相同的性能），而专家层的动量是浪费的。所以 SkewAdam 是用在混合专家模型训练中的精细化优化器，特别适合 GPU 显存有限（如单张 40GB A100/H200）又想训练几十亿参数 MoE 的场景。
文中提到的“分解二阶矩”（Factored Second Moment）是什么意思？为什么能节省那么多内存？传统的 AdamW 优化器对每个参数都保存一个 float32 的二阶矩估计（variance），对于一个 n×m 的权重矩阵，需要存储 n×m 个 float32 值（例如 4096×4096=16.8M 个，约 64MB）。分解二阶矩（来自 Adafactor）的思路是：分别计算每一行的平均值和每一列的平均值作为统计，然后通过外积重建出 rank-one 的二阶矩估计。这样存储量从 n×m 降到了 n+m（例如 4096+4096=8192 个，约 32KB）。对于 128 个专家矩阵，传统 AdamW 需要 128×64MB=8.2GB，而分解版本只需 128×32KB=4.1MB，节省了 2000 倍！但需要注意，这种分解对路由器（gate）不适用，因为路由器中不同 logit 的相对大小对路由决策至关重要，分解会破坏这种信号。
论文中 SkewAdam 的困惑度比 AdamW 还低（108.4 vs 126.8），这是否说明 SkewAdam 本身就是一个更好的优化器？这是一个非常重要的判断。论文自己给出了清晰的回答：分层策略的贡献是内存，而不是更好的优化质量。从表 3 的消融实验可以看到，如果给所有层都加上动量（Uniform with momentum），困惑度也是 108.3，和 SkewAdam 的 108.9 在种子噪声内相同。所以 SkewAdam 之所以困惑度低，是因为它保留了骨干网络的动量（dense梯度），而它的分层分配策略节省了 20 倍的内存。这个差距在扫过学习率后仍然存在（AdamW 调优后 118.5 vs SkewAdam 108.4），说明在 MoE 场景下，均匀的动量分配确实不如骨干网络加分解专家的组合。但你必须清楚：这个好成绩是在 82M token 的小数据规模下获得的，在更大的 token 预算下，AdamW 是否能反超还需要验证。论文的第十行结论说得很谦虚："The contribution is memory, not a better optimizer"，这个定位非常精准。
如果你还有哪些想要了解的，欢迎在评论区留言或者讨论~
龙哥点评
论文创新性分数：★★★★✰
思路新颖实用，把优化器状态按参数角色分层分配，但分解二阶矩和更新clipping来自Adafactor，创新在于分配策略。
实验合理度：★★★★★
实验设计严谨：相同初始化、数据顺序、bfloat16更新路径。设置Adafactor对照、分层消融、学习率扫描、三种GPU复现、零样本评估。不足是大多数实验单种子运行，小规模训练验证空间有限。
学术研究价值：★★★★★
“分层分配”设计原则打破“优化器应对所有参数一视同仁”的隐含假设，对MoE训练、混合精度训练、内存高效优化等领域有启发。
稳定性：★★★★✰
82M token下负载均衡稳定（0.0505，理论下限1%内），但未在多层MoE上验证，更大规模和更长训练表现未知。
适应性以及泛化能力：★★★★✰
主要验证MoE场景（单MoE层），对dense模型、多层MoE、更大模型的适用性不清楚。分层原则泛化可能强，但具体配置需调整。
硬件需求及成本：★★★★★
峰值31.3GB内存（含权重、梯度、激活、状态）在40GB显存上有余量。状态仅1.29GB，节省是结构性的，不增加额外计算开销。
复现难度：★★★★★
代码已开源（MIT协议），单文件实现，无额外依赖。附录提供完整配置表和组件级状态分解，复现条件好。
产品化成熟度：★★★✰✰
集成难度低（改一行优化器名），但只验证82M token单MoE层场景，与产品级百B token、多层MoE差距巨大。零样本评估显示下游能力无竞争力。
可能的问题：论文自信展示困惑度领先，但只经82M token训练，零样本评估接近随机。消融实验单种子运行，无法确认细微差异。逐层策略泛化性未验证。作者建议生产规模需重新验证权重衰减兼容性和bfloat16数值稳定性。
主要参考文献
