行业资讯

独立研究者提出“分层优化器“:让AI训练内存占用缩减97%,普通显卡也能训练百亿参数模型?

发布时间:2026/8/4 1:51:46
独立研究者提出“分层优化器“:让AI训练内存占用缩减97%,普通显卡也能训练百亿参数模型? 这项由独立研究者完成的研究以预印本形式发布于2026年7月21日论文编号为arXiv:2607.19058有兴趣深入了解的读者可以通过该编号在arXiv平台查询完整论文。**一个让AI研究者头疼已久的内存账单**训练大型AI模型就像经营一家大型餐厅。你不仅要备好食材模型参数还要备好厨师的工作记录本优化器状态——记录每道菜上次怎么做、下次应该怎么调整。问题是这本记录本比食材本身还要厚得多。在这项研究中研究者聚焦的是一类叫做稀疏混合专家模型MoE的AI架构。这种架构的设计哲学非常聪明模型里有许多专家处理每条信息时只调用其中一小部分而不是所有专家都同时上阵。这样就做到了参数总量很大但每次实际运算量很小。研究者测试的具体模型有67.84亿个参数但每次处理一个词语时实际激活的参数只有约4.4亿节省了大量计算资源。然而训练时用来教模型学习的优化器——具体来说是业界最常用的AdamW优化器——并不享受这种折扣。AdamW需要给每个参数保存两份学习记录分别叫一阶动量和二阶动量而且是用高精度浮点数存储的。结果就是67.84亿个参数对应的权重数据只占12.6GB但AdamW光是这两份记录就要占50.6GB整个训练峰值内存高达81.4GB——远超市面上主流的40GB专业显卡。这个问题在学术界和工业界都是真实的痛点。很多研究团队为了训练大模型要么租用昂贵的超大内存服务器要么把模型拆分到多张显卡上联合训练成本和复杂度都大幅上升。于是这位独立研究者开始问一个没有人认真追究过的问题**真的有必要对所有参数都保存同样多的学习记录吗****二、混合专家模型里的三类员工**为了理解这项研究的核心洞见我们先要搞清楚这类AI模型的内部结构。可以把这个67.84亿参数的模型想象成一家大型公司里面有三类截然不同的员工群体。第一类是骨干员工也就是论文中所说的密集骨干网络dense backbone包括词嵌入层、注意力机制、密集前馈层等。这批人数量不多只占总参数量的5%约3.41亿但他们是每处理一个词语都要上班的工作量稳定、密集工作记录的意义非常大。第二类是人数最多的专家团队也就是模型里的128个专家网络。这批人占了总参数量的95%约64.43亿是绝对的大头。但他们的工作方式很特殊每次处理一个词语只有其中约2个专家被调用top-2路由机制也就是说每个专家平均只处理1/64的词语。他们的工作是稀疏的、零散的每次上班都是偶发任务。第三类是人数极少但地位关键的调度员也就是路由器router。这个小小的门控网络只有约52万个参数占总量不到0.01%但它的职责是决定每个词语该发给哪些专家处理——它的每一个决策都影响着整个公司的运转效率尤其是各专家之间的工作量是否均衡。以往所有的优化器不管是AdamW、Lion还是Muon都把这三类员工一视同仁给每个人都保存同样格式、同样精度的工作记录。研究者认为这是一种巨大的浪费而且可能并不合理。**三、SkewAdam一套按需分配的记账策略**针对上述问题研究者提出了名为SkewAdam的优化器。这个名字本身就揭示了设计思路Skew的意思是倾斜、偏向也就是把记账资源向真正需要的地方倾斜。对于骨干员工5%的参数SkewAdam保留了完整的float32精度动量缓冲区外加一个分解式二阶动量。动量这个概念可以用物理上的惯性来理解一个滚动的球会保持方向继续滚不会因为某一步地面稍微不平就立刻转向。骨干员工每步都有密集的梯度信号动量能让学习方向更加稳定1.27GB的存储代价在这里是值得的。对于专家团队95%的参数SkewAdam完全不保存动量只保存一个分解式二阶动量。所谓分解式二阶动量是借鉴自Adafactor优化器的技巧对于一个n×m的权重矩阵不直接存储n乘以m个数而是只存储n个行统计量和m个列统计量然后用它们的乘积来近似重建完整矩阵。对于一个4096×4096的专家矩阵这相当于从存储1600万个数减少到只存储8192个数节省了近2000倍空间。不保存动量的原因是每个专家平均每步只接收约128个词语的梯度信号本来就稀疏嘈杂硬要用动量去平滑这样的噪声效果有限代价却高达24GB。对于调度员路由器不到0.01%的参数SkewAdam反而给了最精确的待遇完整的、不分解的float32二阶动量并且不施加权重衰减。原因是路由器的工作性质特殊——它的128个输出数值之间的相对大小直接决定了哪些专家被激活。如果用分解的方式估算就会把不同输出位置的统计量混在一起破坏这种精细的相对关系。2MB的精确存储换来的是整个调度系统的稳定性非常划算。这套策略的整体结果是总优化器状态从AdamW的50.55GB降低到1.29GB只有后者的2.6%。训练峰值内存从81.4GB降到31.3GB终于可以在一张40GB的显卡上运行。**四、具体的数学机制记账本是怎么写的**为了让对技术细节有兴趣的读者有更清晰的认识这里用比较直观的方式解释SkewAdam内部的运作逻辑。分解式二阶动量的更新方式是这样的每一步先分别计算梯度矩阵在行方向和列方向的均方值然后用指数移动平均的方式逐步积累历史信息。重建时把行向量和列向量相乘再除以行向量的均值得到一个近似的完整二阶动量矩阵。这个近似在梯度矩阵的真实二阶动量接近秩为1也就是行列之间高度相关时是精确的在实际训练中通常是足够好的近似。更新步骤还包括一个更新幅度裁剪机制计算出预调整方向后如果它的均方根超过1就把整个向量等比例缩小到均方根恰好等于1。这就好比规定每步最多走一定的距离防止偶发的大梯度导致参数跳飞。主权重是用bfloat16格式存储的一种低精度浮点数每个数只占2字节而float32占4字节。每一步计算完更新量后在将其写回bfloat16权重时会加入一个微小的随机噪声范围是当前权重数值精度的一半这个技巧叫做抖动随机舍入目的是避免系统性的舍入误差积累。值得一提的是这项研究中所有对比的优化器都使用了完全相同的这套写回机制确保比较的公平性。**五、实验设计一场严格的同台竞技**为了让结果有说服力研究者设计了一套极为严格的对照实验。模型固定为同一个67.84亿参数的解码器结构两个模块构成第一个模块用密集SwiGLU前馈层第二个用128专家的MoE层宽度均为4096使用分组查询注意力机制32个查询头8个KV头。训练数据来自OpenWebText语料库按文档哈希值划分为95%训练集和5%验证集确保两侧没有重叠文档。每次运行固定10000步批量大小为64条序列×128个词语总计约8190万个词语单轮不重复。最关键的控制是所有参与对比的优化器都从同一个随机初始化出发按照完全相同的数据顺序进行训练使用相同的bfloat16主权重和相同的写回路径。这就像让四位厨师从同一份食材开始按同一份菜单顺序烹饪只是各自使用不同的烹饪手法——这样最终菜品的差异就只能来自手法本身。参与对比的优化器除SkewAdam外还有AdamW使用3×10??学习率、Lion使用1×10??学习率按其论文建议降低3到10倍、以及Muon矩阵参数学习率0.02其余参数用内置Adam学习率10??。训练在一块英伟达H200显卡141GB显存上进行之所以选择这么大内存的卡是因为要包含AdamW这个对照组——如果换成40GB卡AdamW根本跑不起来。**六、主要结果内存大降的同时收敛更快**实验结果在内存和效果两个维度上都给出了清晰的答案。内存方面SkewAdam的峰值训练内存为31.3GB是唯一能在40GB显卡上运行的方案。AdamW需要81.4GBLion和Muon虽然各自将状态降到25.27GB但由于模型权重、梯度、激活值等其他开销总峰值仍分别达到56.6GB和57.6GB均超过40GB的门槛。吞吐量方面SkewAdam每秒处理约5000个词语比AdamW快6.6%仅比Lion慢1.5%。Muon因为每步都要对64.43亿专家参数运行牛顿-舒尔茨正交化迭代吞吐量只有3409词语/秒比SkewAdam慢了约32%。收敛质量方面训练前3000步AdamW和Muon的验证困惑度衡量语言模型好坏的指标数字越低越好领先于SkewAdam——第1000步时SkewAdam比AdamW落后114个困惑度点。然而从第4000步开始SkewAdam反超最终以108.4的验证困惑度结束训练领先Muon的120.2和AdamW的126.8。Lion的表现非常糟糕最终停在393.7训练过程中困惑度甚至在第9000步之后开始反弹恶化始终未能有效学习。路由均衡方面理想状态下128个专家应该被大致均等地使用此时负载均衡损失函数的理论最低值是0.05。SkewAdam和AdamW从第4000步起就稳定在0.0505和0.0502几乎贴着理论下限。Lion虽然路由相对均衡0.0537但语言模型完全没学好。Muon则在最后1000步突然跳升到0.0608比理论下限高出22%与此同时其验证困惑度也出现了唯一一段平台期两者在时间上高度吻合。**七、与Adafactor和GaLore的额外比较**研究者后续在英伟达H100 NVL 47GB MIG切片上补充了两项对照实验均匀Adafactor和GaLore风格基线。Adafactor是SkewAdam最直接的亲戚——SkewAdam使用的分解式二阶动量估计器就是从Adafactor那里借鉴来的。Adafactor完全不保存动量对所有参数统一使用分解式二阶动量优化器状态只需12MB比SkewAdam的1.29GB还少得多。然而均匀Adafactor的最终验证困惑度是149.5比SkewAdam的109.0足足高出40个点而且在训练最后1000步里改善幅度不到0.1已经近乎停滞。这个对比非常有说服力因为两个优化器共享完全相同的分解式二阶动量机制和更新裁剪策略唯一的区别是SkewAdam给骨干参数保存了动量而Adafactor没有。由此可以清晰地定位SkewAdam之所以效果更好核心原因是动量而不是分层分配策略本身——分层策略带来的是内存节省不是额外的精度提升。GaLore风格基线秩128的低秩梯度投影在这个设置下彻底失败最终困惑度为1839.9路由虽然均衡但语言模型完全没有学到有效表示。研究者明确指出这不能作为对GaLore方法本身的否定——他们使用的是自己的实现版本未经系统调参与GaLore原始论文的设置有所不同。但这个失败案例至少提示把稀疏专家梯度投影到低秩子空间这件事需要非常谨慎。**八、哪个分层决策真正起作用消融实验的答案**一个合理的质疑是SkewAdam的效果是因为分层策略还是因为其中某个特定的设计选择研究者在AMD MI300X192GB上运行了消融实验系统地把各个决策一一关闭观察结果变化。实验包含四个变体完整的SkewAdam策略在专家上恢复动量的版本把路由器二阶动量改为分解式的版本以及完全均匀分配所有参数都用动量加分解式二阶动量的版本。结果是四个变体的验证困惑度全部落在108.2到108.9之间互相之间的差异完全在单次运行的随机噪声范围内。路由均衡损失也全部在0.050附近几乎没有区别。真正发生显著变化的只有优化器状态大小完整SkewAdam是1.29GB恢复专家动量后变为25.29GB完全均匀版本同样是25.29GB。这个结果说明了两件事。其一分层策略的贡献是内存而不是额外的困惑度提升——完整策略和均匀策略达到完全相同的训练效果前者只用了后者二十分之一的优化器状态。其二专家层的动量是纯粹的浪费加上它花费24GB但困惑度只变动了0.2完全在噪声范围内。路由器使用精确二阶动量的选择而非分解式也被证明既无害也无益——改为分解式后困惑度和路由均衡均无变化。研究者将其描述为无害但非关键的选择保留它只是出于对路由器特殊角色的谨慎考量。**九、调参之后基线能追上吗**为了排除SkewAdam用了更合适的学习率这一可能的解释研究者系统地对AdamW和Adafactor进行了学习率扫描而SkewAdam保持不调参的默认设置3×10??。AdamW扫描了10??、3×10??、10??三个学习率每个跑三个随机种子。最优结果在10??时出现达到118.5±0.5比默认的126.8改善了约8个点。Adafactor扫描了3×10??到3×10??共五个学习率最优结果同样在10??达到139.7两个种子平均差异0.005。值得注意的是3×10??时Adafactor出现了欠训练困惑度257.5而更高的学习率全部更差说明10??的选择已经是有效区间内的最优点上下都有了明确的边界。调参之后差距缩小了但没有消弭未经调参的SkewAdam108.4到109.0之间领先最优AdamW约10个困惑度点领先最优Adafactor约30个点。10个点的差距约是AdamW种子间标准差的20倍统计上非常显著。研究者还指出调过参的AdamW和调过参的Adafactor之间仍然差了约21个点118.5对139.7而这两者的唯一区别正是动量——AdamW有动量Adafactor没有。这个跨优化器的对比与消融实验的结论完全吻合动量是关键分层分配是内存的解法。**十、零样本评估诚实的不知道**研究者还用标准的语言模型评估工具对训练结束的四个模型进行了零样本推理测试任务包括PIQA物理常识推理、WinoGrande代词消歧、HellaSwag故事续写和ARC-Challenge科学问答。结果很诚实在仅8190万词语的训练量之后所有四个模型在这些任务上的表现几乎全部接近随机猜测水平。PIQA上四者得分在53.5%到55.9%之间随机猜测是50%HellaSwag和ARC-Challenge接近随机猜测的25%。任何两个优化器之间的差异都在一到两倍标准误差之内没有统计显著性。8190万词语对于现代大语言模型来说只是零头——正式训练通常需要数千亿甚至数万亿词语。研究者坦然地将这些零样本结果定性为完整性检查而非能力证明并明确提醒读者困惑度上的提升不等于下游任务的提升不要对这份结果赋予超出它本身的意义。**十一、局限与未竟之处**研究者在论文中对这项工作的局限性进行了相当坦诚的讨论。模型只有两个模块这是出于成本考量的主动选择——它把95%的参数集中在一个专家库里制造了一个极端的压力测试场景。但真实的MoE模型通常有几十乃至上百个交替的注意力层和MoE层分层策略在多层叠加的复杂路由结构下是否同样有效尚未得到验证。大多数配置只跑了一个种子消融实验的四个变体在0.6个困惑度点的范围内波动而这恰好是单次运行间随机噪声的量级因此消融结果应该被理解为内存代价相当效果相当而不是对各变体之间存在任何具体顺序的声明。训练只有8190万词语上下文长度只有128个词语按现代标准极为短小。权重衰减在bfloat16精度下是一个无效操作变化量远小于浮点精度被直接舍入为零这意味着整个实验实际上是在没有权重衰减的条件下进行的与真实的长时间生产训练有所不同。如果要将这套方案应用于生产规模必须把权重衰减项整合进float32更新步骤在精度截断之前完成计算。GaLore风格基线使用的是研究者自己写的实现未经系统调参结果不代表GaLore方法的真实上限。Lion和Muon也只测试了单个学习率没有进行系统调参其表现可能还有提升空间。**归根结底这项研究告诉了我们什么**说到底这项工作的核心贡献是一个朴素却被忽视的观察混合专家模型里的三类参数工作方式截然不同理应得到不同的记账待遇。把骨干网络的动量、专家层的分解式方差估计、和路由器的精确方差估计组合在一起优化器状态从50.55GB降到1.29GB训练峰值内存从81.4GB降到31.3GB而最终模型效果不但没有变差反而因为保留了骨干层的动量而优于AdamW。这对于想在资源有限的条件下训练大模型的研究者和工程师来说是一个实际可用的工程方案。它不依赖量化、不依赖梯度压缩、不依赖多卡分布式只是重新思考了每一分内存都值得花在哪里这个问题。研究者自己也强调这不是一个更好的优化算法而是一个更好的内存分配策略——Adam家族的效果花了2.6%的Adam状态存储。当然在更深的模型、更长的训练、更多的随机种子上验证这套方案是这项工作留给未来的任务。但作为一个方向它提出了一个颇具启发性的设计原则在构建优化器时先问清楚这部分参数是谁它们是怎么工作的再决定给它们记多少账。有兴趣深入了解的读者可以通过arXiv:2607.19058查阅完整论文代码也已公开在GitHub上供复现使用。---QAQ1SkewAdam和AdamW的训练效果哪个更好A在这项研究的实验设置下SkewAdam的最终验证困惑度为108.4优于AdamW的126.8。即使对AdamW进行学习率调参其最优结果也只能达到118.5仍然落后于未调参的SkewAdam约10个点。但需要注意实验规模较小约8190万词语且模型结构特殊结果能否推广到更大规模仍需验证。Q2SkewAdam为什么不给专家层保存动量A因为每个专家平均每步只处理约128个词语的梯度信号稀疏且噪声高动量平滑的效果有限却要消耗24GB内存。消融实验证实给专家层加回动量后验证困惑度变化仅0.2个点完全在随机噪声范围内属于纯粹的浪费。Q3分解式二阶动量是怎么节省内存的A对于一个n×m的权重矩阵传统方法要存储n×m个数分解式方法只需存储n个行统计量和m个列统计量通过它们的乘积来近似完整矩阵。以4096×4096的专家矩阵为例存储从1600万个数减少到约8192个数节省约2000倍整个专家库的二阶动量只需12MB。