训练大模型像收拾行李:模型权重只是衣服,负责更新权重的“优化器状态”往往才是最占地方的箱子。Mixture-of-Experts(MoE,混合专家模型)虽然每次只调用少数“专家”来控制计算量,全部专家及其训练状态却仍要装进显存。SkewAdam想解决的,就是这笔容易把单张GPU撑爆的额外开销。
它没有平等对待所有参数,而是按参数行为分层保存状态:占参数约5%的骨干部分保留动量和压缩后的二阶统计量;占约95%的专家只保留后者;不足0.01%的路由器——负责决定调用哪些专家——则保留精确统计量。数值精度越低越省显存,但误差风险也更高,这套设计把有限空间留给更敏感的部分。
据作者在 Reddit 的介绍,SkewAdam把一个6.78B参数MoE的优化器状态显存从50.6GB降至1.29GB,减少97.4%;训练峰值显存从81.4GB降至31.3GB,因此可装进单张40GB GPU。作者还称收敛和路由稳定性没有受损,但供稿未提供独立复现结果,这条低成本训练路径仍值得验证。