---
title: "优化器状态应放在哪里？面向内存高效MoE训练的分层状态分配"
source: "arXiv AI/CL/LG（论文）"
source_tier: "t2"
url: "https://arxiv.org/abs/2607.19058v1"
permalink: "https://feed.astralor.com/items/itm_b4bae957f7521b6e57c79ab6"
published: "2026-07-21T12:47:15.000Z"
score: 83
channel: "ai"
tags: ["AdamW","SkewAdam","分层状态分配","动量优化","因子化二阶矩","混合专家模型 (MoE)"]
---

# 优化器状态应放在哪里？面向内存高效MoE训练的分层状态分配

## AI 点评

> 这篇论文对训练MoE模型的团队非常有价值：它证明了通过精细的优化器状态分配，可以在不牺牲精度的情况下大幅降低显存占用，甚至还能提升性能。尤其适合那些受限于40GB或80GB GPU显存的实验室或企业。值得仔细阅读其分层策略和动量保留的消融分析，可为实际训练框架（如DeepSpeed、Megatron）提供优化方向。

## 速览

SkewAdam通过为MoE的骨干、专家和路由器分配不同精度的优化器状态，将AdamW需要的50.6GB状态压缩至1.29GB（2.6%），峰值训练内存从81.4GB降至31.3GB，且验证困惑度（108.4）优于AdamW（126.8）、Muon（120.2）和Lion（393.7）。实验表明，状态分配的位置与数量同样重要，保留动量是精度关键，而分层分配在不损失精度前提下大幅节省内存。

## 关键点

- 优化器状态是MoE训练中最大的内存开销，占AdamW内存的80%以上。
- SkewAdam对骨干（5%参数）使用float32动量+因子化二阶矩，专家（95%）仅因子化二阶矩，路由器（<0.01%）使用精确二阶矩。
- 相比AdamW，状态内存减少97.4%，峰值内存降低61.5%，且困惑度提升。
- 动量是保持精度的关键，单纯减少状态（如Adafactor）会导致性能大幅下降。

## 全文

Optimizer state is the largest single line item in the memory budget of mixture-of-experts (MoE) training: on a 6.78B-parameter MoE language model, AdamW keeps 50.6 GB of first and second moments to update 12.6 GB of bfloat16 weights. We study SkewAdam, an optimizer built on the observation that the three parameter populations of an MoE - the dense backbone, the experts, and the router - differ enough in size and gradient statistics that they should not receive the same state. SkewAdam keeps float32 momentum plus a factored second moment for the backbone (5% of parameters), a factored second moment alone for the experts (95%), and an exact second moment for the router (<0.01%). The resulting state occupies 1.29 GB, 2.6% of AdamW's, and peak training memory falls from 81.4 GB to 31.3 GB, within the budget of a 40 GB accelerator. In a controlled comparison from identical initializations over 82M tokens, SkewAdam reaches validation perplexity 108.4, ahead of AdamW (126.8), Muon (120.2), and Lion (393.7), and settles router load balance to within 1% of its uniform floor. The allocation is not what earns that perplexity: a tier ablation matches it with twenty times the state, and Adafactor, which shares the factored estimator but drops momentum, plateaus 40 points behind. The tiers buy memory at no cost to accuracy; the accuracy comes from keeping momentum, which a uniform optimizer shares too. Sweeping the baselines' learning rates narrows but does not close the gap: the best tuned AdamW reaches 118.5, tuned Adafactor 139.7. Where optimizer state lives, these results suggest, matters at least as much as how much of it there is.

---

引用前请核对原文;本文由星际流动整理:https://feed.astralor.com/items/itm_b4bae957f7521b6e57c79ab6