AI 快讯SkewAdam把MoE状态内存砍掉97%
开发心得

SkewAdam把MoE状态内存砍掉97%

2026-07-22T10:06:45.675Z
SkewAdam把MoE状态内存砍掉97%

SkewAdam 通过按参数角色分层分配优化器状态,将 MoE 训练的状态内存从 50.6 GB 降至 1.29 GB,并让 6.78B 参数模型装进单张 40GB GPU。

SkewAdam把MoE状态内存砍掉97%

**SkewAdam 已于近日开源,它把一个 6.78B 参数 MoE 模型的优化器状态内存从 50.6 GB 压缩到 1.29 GB,降幅达到 97.4%。**截至 2026 年 7 月 22 日,项目作者已经公开论文预印本与 GitHub 代码;在其硬件测试中,训练峰值显存也从 81.4 GB 降至 31.3 GB,使这一级别的 MoE 模型能够在单张 40GB GPU 上训练。

这不是又一个把 Adam 状态改成 8-bit 的常规方案。SkewAdam 的关键判断是:**MoE 不同模块的参数行为并不相同,因此没有必要为每个参数分配完全相同的优化器状态。**它把主干网络、专家网络和路由器拆成三个层级,分别保留不同精度、不同形态的一阶与二阶统计量。

SkewAdam将MoE参数划分为Backbone、Experts和Router三层,并为其分配不同优化器状态的架构示意图

真正吃掉显存的,往往不是模型权重

**优化器状态是训练过程中用于记录参数更新历史的额外张量,例如 AdamW 的一阶动量和二阶矩估计。**推理时只需要保存模型权重和少量运行时缓存,训练时却还要保存梯度、激活值以及优化器状态,因此“模型能推理”与“模型能训练”是两套完全不同的显存账本。

**AdamW 的状态开销通常比一份 BF16 模型权重高出数倍。**如果一阶动量 m 和二阶矩 v 都以 FP32 保存,每个参数仅优化器状态就需要 8 字节;相比之下,一份 BF16 权重通常只需要每参数 2 字节。换句话说,在不考虑梯度、主权重副本和激活值的情况下,AdamW 状态就可能达到 BF16 权重体积的 4 倍。

**SkewAdam 论文给出的测试结果正好暴露了这笔成本。**被测试模型本体约占 12.6 GB,但 AdamW 优化器状态达到 50.6 GB,优化器本身才是显存预算中最大的一项。加上权重、梯度、激活值和临时缓冲区后,训练峰值达到 81.4 GB,单张 80GB GPU 都可能处在溢出边缘,更不用说常见的 24GB、40GB 或 48GB 显卡。

| 显存项目 | AdamW 基线 | SkewAdam | 变化 | |---|---:|---:|---:| | 模型规模 | 6.78B MoE | 6.78B MoE | 不变 | | 模型权重体积 | 约 12.6 GB | 约 12.6 GB | 不变 | | 优化器状态 | 50.6 GB | 1.29 GB | 减少 49.31 GB,降幅 97.4% | | 训练峰值显存 | 81.4 GB | 31.3 GB | 减少 50.1 GB,降幅 61.5% | | 最低可用硬件级别 | 超过单卡 80GB 的安全范围 | 单张 40GB GPU | 单卡可运行 |

**“状态内存减少 97.4%”不等于“总训练显存减少 97.4%”。**模型权重、梯度和激活值不会因为更换优化器而消失,所以最终峰值显存从 81.4 GB 降到 31.3 GB,对应的整体降幅是 61.5%。这已经足够改变硬件门槛,但在传播这些数字时必须区分优化器状态与训练总显存。

SkewAdam为什么对MoE特别有效

**混合专家模型(Mixture-of-Experts,MoE)是一类通过路由器为每个输入只激活部分专家网络的稀疏模型。**MoE 可以在不按比例增加单次计算量的情况下扩充总参数量,但“本次前向没有激活某个专家”并不意味着训练框架无需为该专家保存参数和优化器状态。

**MoE 的参数分布天然高度倾斜,这正是 SkewAdam 名称中“Skew”的来源。**在作者测试的结构里,专家参数约占总参数的 95%,主干网络约占 5%,路由器则不到 0.01%。传统 AdamW 把三类参数一视同仁,为全部参数各保存一份完整的一阶动量和二阶矩;SkewAdam 则利用这种比例失衡,把最昂贵的状态只留给少数真正需要它的模块。

| 参数层级 | 参数占比 | SkewAdam 保存的状态 | 设计意图 | |---|---:|---|---| | Backbone 主干网络 | 约 5% | 一阶动量 + 因式分解二阶矩 | 保留较稳定的更新方向,同时压缩二阶状态 | | Experts 专家网络 | 约 95% | 仅因式分解二阶矩 | 移除覆盖绝大多数参数的一阶动量,获得主要内存收益 | | Router 路由器 | 小于 0.01% | 精确二阶矩 | 路由参数极少,可以用很低成本保留细粒度统计 |

**主干网络是被所有 token 共享的公共路径,因此 SkewAdam 为它保留了一阶动量。**一阶动量可以理解为参数近期更新方向的滑动平均,它能抑制梯度抖动并形成更连续的更新轨迹。主干只占约 5% 参数,即使继续保存这部分状态,成本也远低于为全部专家保存完整动量。

**专家网络是状态压缩的主战场,因为它们占据约 95% 的参数。**SkewAdam 不再为专家保存完整的一阶动量,只保留因式分解后的二阶矩。仅这一项取舍,就避开了按专家参数量线性增长的一大块 FP32 状态。

**路由器是 MoE 中参数极少但训练敏感度很高的模块。**路由器决定 token 被送往哪些专家,微小偏差可能导致负载失衡、部分专家过载或专家利用率下降。SkewAdam 因此没有继续压缩这部分状态,而是保留精确二阶矩;由于路由参数不足总量的 0.01%,这笔“精打细算后的奢侈”几乎不影响总显存。

因式分解二阶矩,省的是矩阵级状态

**因式分解二阶矩是一种用行统计量和列统计量近似完整二阶矩矩阵的方法。**对于形状为 n × m 的权重矩阵,完整二阶矩需要保存 n × m 个数,而因式分解后通常只需保存长度为 nm 的两个统计向量,存储复杂度从 O(nm) 降到 O(n+m)

**这种压缩对大尺寸线性层尤其有效。**假设一个专家矩阵大小为 4096 × 14336,完整 FP32 二阶矩需要保存约 5872 万个数,单个矩阵约占 224 MiB;若只保存对应的行、列统计量,则只需 18432 个 FP32 数值,约 72 KiB。实际优化器还会涉及分块方式、缩放与数值稳定处理,但数量级差异已经非常明显。

**SkewAdam 并不是第一个采用因式分解状态的优化器,但它把这种方法与 MoE 的模块分工结合了起来。**Adafactor 已经证明二阶矩不一定要逐元素完整保存,8-bit Adam 则通过低位量化压缩状态;SkewAdam 更进一步,先问“这类参数是否需要一阶动量”,再问“二阶矩是否需要完整保存”。它的收益来自状态种类、状态形态与参数角色三方面同时裁剪。

| 方案 | 主要压缩方式 | 是否针对 MoE 参数结构 | 主要代价或限制 | |---|---|---:|---| | AdamW | 不压缩,保存完整一阶与二阶状态 | 否 | 状态内存最高,但行为成熟、生态完善 | | 8-bit Adam | 将优化器状态量化到低位宽 | 否 | 仍为大部分参数保留两类状态,需要量化与缩放逻辑 | | Adafactor 类方案 | 因式分解二阶矩,通常减少或省略一阶状态 | 否 | 对参数形状和训练配方更敏感 | | ZeRO / FSDP | 把参数、梯度或状态切分到多张 GPU | 否 | 降低单卡占用,但引入通信和多卡依赖 | | SkewAdam | 按主干、专家、路由器分层分配状态 | 是 | 目前验证范围有限,需要更多模型与训练规模复现 |

**SkewAdam 与 ZeRO、FSDP 并不是非此即彼的关系。**前者减少优化器状态的总量,后两者主要把现有状态分散到多张设备;理论上两类方法可以叠加。对单卡实验而言,SkewAdam 直接降低了准入门槛;对多卡训练而言,它可能进一步减少显存占用和状态搬运量,但通信收益仍需真实集群测试确认。

1.29 GB这个数字是怎么来的

**SkewAdam 的 1.29 GB 状态占用,与“只给约 5% 参数保留动量”的设计在数量级上基本一致。**如果 AdamW 基线对应约 50.6 GB 的两份 FP32 状态,那么单独一份完整状态约为 25.3 GB;其中 5% 约为 1.27 GB,已经非常接近论文报告的 1.29 GB。因式分解二阶矩和极小规模的路由器状态,只在这个基础上增加有限开销。

**这组数据也说明真正决定结果的不是单纯降低精度,而是删除不必要的逐参数状态。**把 50.6 GB 的状态量化到 8-bit,理论体积仍会受到两份逐参数数组的约束;SkewAdam 对 95% 的专家参数直接取消一阶动量,并将二阶矩从矩阵级数组改成行列统计,其压缩上限自然更高。

**模型权重体积与总参数量之间存在少量口径差异并不罕见。**6.78B 参数若全部按 BF16 直接计算,原始字节数会略高于 12.6 GB;论文中的数字可能受到 GB 与 GiB 单位、非训练参数、统计范围或具体模型实现影响。因此,最值得关注的是同一测试口径下 50.6 GB 到 1.29 GB、81.4 GB 到 31.3 GB 的对照,而不是跨框架机械换算后的个位数差异。

单张40GB GPU训练6.78B MoE,意义不只在省卡

**SkewAdam 最直接的价值是把原本需要高显存卡或多卡切分的实验,压回单张 40GB GPU。**31.3 GB 的峰值占用还留下约 8.7 GB 余量,可用于框架波动、数据批次变化和部分缓存,不过实际可用空间仍取决于 CUDA 上下文、显存碎片、序列长度与激活检查点配置。

**单卡可训练会显著降低 MoE 研究的工程复杂度。**开发者不必为了验证一个路由策略就先搭建分布式环境,也不必处理参数分片、跨卡同步、通信超时和拓扑差异。对于高校实验室、小团队和个人研究者,这种变化比单纯缩短若干百分点训练时间更实用。

**更低的状态内存还可能允许研究者把显存重新分配给更长上下文或更大批次。**MoE 训练中的路由稳定性、专家负载均衡和梯度方差都可能受 batch 组成影响,如果节省出的显存能换来更合理的 token batch,优化器压缩可能间接改善训练效率。不过,这属于可验证的潜在收益,不能直接从当前 31.3 GB 的结果推导出更好的最终精度。

最大疑问:去掉专家动量会不会伤收敛

**SkewAdam 最需要进一步验证的地方,是专家参数在长期训练中是否真的不需要一阶动量。**作者报告称该方案在现有实验中没有牺牲收敛表现,但预印本与单一代码仓库尚不足以证明它能无损替代所有 MoE 训练配方。

**专家梯度具有稀疏、间歇和分布不均的特点,这既是删除动量的理由,也是潜在风险。**某个专家只接收部分 token,其梯度更新不像稠密主干那样连续;历史动量可能价值较低,也可能在专家长时间未被选中后变成陈旧信号。另一方面,当数据分布变化或路由发生漂移时,缺少动量也可能使专家更新更依赖当前 batch,增加训练噪声。

**真正有说服力的后续验证至少需要覆盖四类变量。**第一是从 6.78B 扩展到数十亿乃至百亿参数以上;第二是增加专家数量和改变 top-k 路由策略;第三是拉长训练 token 数,观察后半程损失曲线;第四是检查路由负载、专家利用率、下游能力和最终困惑度,而不只是确认训练没有发散。

**优化器状态变小也不代表训练一定更快。**因式分解统计需要额外归约操作,分层更新可能带来更多小算子与参数分组;如果实现无法良好融合,单步耗时甚至可能上升。反过来,更少的显存读写也可能缓解带宽压力,因此最终吞吐量必须通过 tokens/s、step time 和 GPU 利用率实测,而不能只看显存数字。

现阶段值得用,但更适合研究和验证

**SkewAdam 已经值得 MoE 开发者试用,因为它解决的是一个真实且昂贵的瓶颈。**从 50.6 GB 到 1.29 GB 的优化器状态缩减不是边角改进,而是足以改变硬件选择的数量级变化;它也给出了一个比“统一量化全部状态”更贴合 MoE 结构的方向。

**SkewAdam 暂时还不适合被直接视为 AdamW 的通用替代品。**当前结果主要证明了分层状态分配在特定 6.78B MoE 设置中的可行性,尚未形成跨架构、跨数据集、跨训练周期的大规模复现。生产训练在切换前,应至少对齐损失曲线、最终指标、路由负载、训练吞吐量和断点恢复行为。

开发者评估 SkewAdam 时应优先记录五项指标:

  1. 峰值显存与稳定训练所需的安全余量;
  2. 每秒处理 token 数和单步耗时;
  3. 与 AdamW 对齐步数下的训练损失;
  4. 各专家接收 token 的比例与负载均衡情况;
  5. 长周期训练后的下游指标、困惑度和数值稳定性。

**SkewAdam 最有价值的启发不是“所有模型都该删除动量”,而是优化器状态不必继续采用一刀切设计。**MoE 本身已经用条件计算区分不同参数的参与方式,优化器却长期按稠密模型思路为所有参数准备同样的历史状态。SkewAdam 把条件计算继续推进到优化器层面:常用的主干保留动量,占比最大的专家使用低成本统计,敏感但极小的路由器获得精确处理。

**如果后续独立复现能够确认收敛质量,SkewAdam 可能成为单卡 MoE 训练中非常实用的一块基础设施。**它未必取代 AdamW,也未必适用于每一种模型,但“6.78B MoE、31.3 GB 峰值显存、单张 40GB GPU 可训练”已经是一组足够吸引开发者动手验证的数据。

参考来源

相关推荐

查看全部