最佳化器狀態
3 秒看懂
- 最佳化器狀態是梯度下降過程中輔助更新模型權重的記憶變數，典型形式為一階動量（速度）與二階動量（梯度平方的指數移動平均），少數高階演算法還包含預條件矩陣等。
- 它決定了每一步引數調整的方向與步長的“動量 + 自適應縮放”，直接影響收斂速度、訓練的穩定性以及最終模型的泛化效能。
- 在大規模訓練中，最佳化器狀態佔用的視訊記憶體經常遠超模型權重本身，是千億引數模型分散式訓練的核心記憶體瓶頸；降低狀態開銷已成為系統研究和工程最佳化的主戰場。
3 分鐘產業解釋
在神經網路訓練中，僅依靠當前這一步的梯度來更新引數往往既低效又不穩定——損失曲面充滿噪聲和狹長峽谷。最佳化器狀態就是最佳化演算法為了記住“過去的梯度經歷”而保持的一組額外變數，比如動量緩衝、梯度平方的滑動平均等。這些狀態幫助演算法像經驗豐富的司機那樣，記住之前的道路傾斜度和顛簸感，從而在一致的方向上加速，在抖動劇烈的方向上減速，大幅度提升收斂效率。
以最常用的 Adam 最佳化器為例，它為每個可訓練引數維護兩個 32 位浮點狀態：一階矩 m（動量）和二階矩 v（梯度平方的無偏估計）。假設一個大型語言模型擁有 100 億個引數，那麼狀態數量就是 100 億 × 2 = 200 億個。在標準混合精度訓練方案中，引數和梯度採用 16 位浮點，而最佳化器狀態保持 32 位浮點以保證數值穩定性，此時僅最佳化器狀態就需要 200 億 × 4 位元組 = 80 GB 的視訊記憶體。再加上 fp16 引數、梯度和一份 fp32 的主權重副本，總視訊記憶體可輕鬆超過 160 GB——這遠遠超出了單張 A100 (80 GB) 甚至 H100 (80 GB) 的容量。正是這種成倍的儲存需求，使最佳化器狀態成為大型模型訓練的“視訊記憶體黑洞”。
產業界解決這一難題的路線已經非常清晰：ZeRO 最佳化器（源自微軟 DeepSpeed）將狀態切分成 N 份分佈到 N 個 GPU 上，每張卡只需儲存 1/N 的狀態，以集合通訊換取視訊記憶體；8-bit 最佳化器（如 bitsandbytes）使用塊級量化將 32 位狀態壓縮為 8 位，記憶體直接降至 1/4；零階最佳化器（如 MeZO）更激進地完全拋棄狀態，僅通過兩次前向傳播的隨機擾動來估計梯度，使單卡微調千億模型成為可能。這些技術已從實驗室快速進入工業界，支撐起了 GPT-4、Llama 3 等超大型模型的訓練與微調。
技術原理
形式化定義與狀態演化
給定引數 θ，損失函式 L(θ)，在迭代步 t 的梯度為 g_t = ∇L(θ_t)。最佳化器狀態 s_t 根據事先定義的規則更新：
s_{t+1} = f(s_t, g_t)而後引數被更新：
θ_{t+1} = θ_t - η·u(g_t, s_t)其中 η 為學習率，u(·) 是綜合利用當前梯度和歷史狀態生成更新向量的函式。不同最佳化器的差異主要體現在 f 和 u 的設計上。
最常見的狀態維度：
- 一階動量 m_t：對歷史梯度做指數滑動平均（EMA），代表梯度的長期趨勢。典型更新式為
m_t = β_1 m_{t-1} + (1-β_1) g_t。 - 二階動量 v_t：對梯度平方做 EMA，反映每個引數的近期波動幅度，用於逐引數縮放學習率。典型式為
v_t = β_2 v_{t-1} + (1-β_2) g_t^2。 - 高階狀態：如 K-FAC 中的 Fisher 資訊矩陣的逆近似、Shampoo 中的預條件矩陣等，形狀為 d×d，儲存開銷 O(d²)，僅限於小規模實驗。
關鍵演算法的狀態構成與記憶體倍數
假設模型總引數量為 P（個數），狀態與引數的儲存格式均為 32 位浮點（4 位元組）。下表給出各經典最佳化器需要維護的狀態變數及其相對於引數量的儲存倍數。
| 最佳化器 | 狀態變數 | 狀態數 | 相對於引數量 P 的儲存倍數 | 備註 |
|---|---|---|---|---|
| SGD（純無狀態） | 無 | 0 | 0× | 不攜帶任何歷史資訊 |
| SGD + Momentum | 速度緩衝 v | 1×P | 1× | 與引數量完全相同 |
| Adam / AdamW | m, v | 2×P | 2× | 工業化訓練中額外記憶體佔比最大 |
| AdaGrad | 累計梯度平方 G | 1×P | 1× | 單調遞增導致後期步長趨近於零 |
| RMSProp | 梯度平方的滑動平均 v | 1×P | 1× | 通常與動量搭配，未單獨列出 |
| LAMB / LARS | 同基礎演算法 + 層自適應縮放因子 | 2×P + 少量統計量 | ~2× | 大 batch 訓練常用 |
| 8-bit Adam（bitsandbytes） | 量化後的 m, v + 量化常數 | 2×P (8-bit儲存) | 0.5×* | 以 8-bit 儲存，相比 fp32 狀態的位元組量降至 1/4 |
| MeZO（零階） | 無 | 0×P | 0× | 基於 SPSA 梯度估計，完全免狀態 |
*注：0.5× 指的是相對於 fp32 狀態（2×P×4位元組）而言，8-bit 狀態僅需 2×P×1 位元組，即 0.5×P×4 位元組的等效倍數，為方便與上表中“儲存倍數”直接對比而列。在實際混合精度環境中，8-bit 狀態的絕對位元組數約為引數量 fp16 引數的同等數量級。
記憶體牆與分散式分片
大型模型混合精度訓練（引數、梯度 fp16，最佳化器狀態 fp32，另存 fp32 主引數）的單卡視訊記憶體構成可以量化為：
- 引數（fp16）: 2×P 位元組
- 梯度（fp16）: 2×P 位元組
- 最佳化器狀態（fp32）: 對於 Adam，m 與 v 各 4×P 位元組，共 8×P 位元組
- 主權重副本（fp32）: 4×P 位元組 合計 ≈ 16×P 位元組。若一個 175B 模型（P=1.75×10¹¹），理論單卡需求高達 16 × 175 × 10⁹ ≈ 2800 GB，顯然無法單卡承載。
ZeRO-1（最佳化器狀態分片） 將 m、v 按資料並行維度切分到 N 張 GPU 上，每張卡僅持有 1/N 的狀態，狀態儲存降為 (8×P)/N。當 N 較大時，單卡狀態視訊記憶體顯著降低。但引數、梯度及主權重副本仍為全量保留，因此最低單卡總視訊記憶體約 8×P（在 N→∞ 時）。若需進一步壓縮，分散式架構還會採用 ZeRO-2（梯度分片） 和 ZeRO-3（引數分片），最終可將每卡模型狀態開銷逼近 16×P / N。
為直觀理解，下方 ASCII 示意圖展示了帶 Adam 狀態的單步訓練流程：
前向 + 反向 (fp16)
模型引數 θ_fp16 → 前向 → 損失 loss
→ 反向 → 梯度 g_fp16
最佳化器步驟（單卡，混合精度）:
梯度 g_fp16 轉換至 fp32
m_fp32 = β1·m_fp32 + (1-β1)·g_fp32
v_fp32 = β2·v_fp32 + (1-β2)·g²_fp32
θ_fp32 = θ_fp32 - η · m_hat / (√v_hat + ε)
將更新後的 θ_fp32 轉為 θ_fp16 供下一步前向
(分散式 ZeRO-1: 在各卡區域性計算梯度後，先對梯度 AllReduce，然後每張卡僅更新自己負責的那部分狀態和引數片段)此外，DeepSpeed 的 ZeRO-Infinity 進一步支援將最佳化器狀態解除安裝至 CPU 記憶體或 NVMe 固態儲存，形成 GPU-CPU-NVMe 三級狀態快取，以頻寬換容量。
關鍵引數
理解和調優最佳化器狀態相關的核心引數，是訓練工程師的基本功。以下列出影響狀態行為與資源消耗的關鍵項。
- β₁（一階矩衰減係數）：控制動量對歷史梯度的記憶長度，通常取 0.9。值越接近 1，記憶越久遠，更新越平滑，但在損失曲面劇烈變化時反應滯後。
- β₂（二階矩衰減係數）：梯度平方的滑窗引數，典型值為 0.999。它決定了自適應學習率對近期梯度波動的敏感度，過高會使適應性不足，過低會引入噪聲。
- ε（數值穩定常數）：作用在分母 √v + ε 中，常取 1e-8 或 1e-7，防止除零。在混合精度下，某些方案會適當增大 ε 以避免半精度下溢。
- 狀態精度：fp32 最穩定，但佔用 4 位元組/狀態；使用 bf16 或 fp16 可能丟失精度，導致訓練發散；8-bit 量化狀態（如 bitsandbytes）通過分塊尺度因子在 1 位元組/狀態下實現收斂。
- 狀態分片數（ZeRO stage）：在多卡訓練中，通過配置 DeepSpeed ZeRO stage 1/2/3 調整狀態、梯度和引數的分佈程度。分片越徹底，單卡記憶體越低，但通訊量也會增加。
- 解除安裝策略：是否將狀態解除安裝到 CPU/NVMe；啟用後狀態不在 GPU 常駐，每步需 PCIe 搬運，訓練速度下降，但可訓練更大型模型。
- 狀態更新頻率：每次迭代都更新；但某些記憶體最佳化策略（如梯度累積）會每 K 步才執行一次最佳化器步驟，期間狀態保持不動，可調增吞吐但不改變狀態峰值儲存。
上述引數共同定義了“時間-空間-精度”的三元權衡，是大型模型訓練資源規劃的核心輸入。
技術路線
按狀態設計模式的差異，可將現有最佳化器劃分為幾條典型路線。
| 路線類別 | 代表演算法 | 狀態數量級 | 記憶體壓力 | 收斂效率 | 大型模型適配度 | 當前階段 |
|---|---|---|---|---|---|---|
| 無狀態 | SGD | 0 | 無 | 低，調參苛刻 | 不適於千億級 | 基線，簡單任務中仍用 |
| 一階動量 | SGD+Momentum | O(P) | 低 | 中 | 中等 | 經典 CV 基線 |
| 自適應（一階+二階） | Adam, AdamW, LAMB | O(2P) | 高 | 高效，魯棒 | 需記憶體最佳化（分片/量化） | 主流標配，生產級訓練基石 |
| 高階預條件 | K-FAC, Shampoo | O(d²) ~ O(P²) | 極高 | 單步改進大 | 不可行 | 研究階段，限於小引數量場景 |
| 零狀態/零階 | MeZO | 0 | 無 | 收斂慢 | 極好，單卡微調巨量模型 | 新興方向，記憶體受限微調場景專用 |
在實際工業部署中，自適應路線（Adam 家族） 仍佔絕對主導，但已不是樸素的單機 fp32 版本。產業界通過 分片（ZeRO）、量化（8-bit/FP8） 和 解除安裝 等手段，將其記憶體壓力大幅稀釋，從而繼續享受自適應狀態帶來的收斂紅利。零狀態路線（MeZO 等）則在邊緣微調和隱私保護領域展現出獨特價值，與自適應路線形成互補。
上游
最佳化器狀態技術的上游主要包括：
- 自動微分架構：PyTorch、JAX、TensorFlow 提供反向傳播計算出精確梯度 g_t，是狀態的“原料”。架構對自定義狀態更新（如
torch.optim.Optimizer）的開放程度直接影響新型狀態設計的實現難度。 - 數值精度生態：NVIDIA GPU 的 Tensor Core 對 fp16/bf16/fp8 的加速，以及 AMD、Intel 的 AI 加速器支援的浮點格式，決定了不同精度下狀態計算的吞吐與能耗。例如 H100 的 FP8 支援推動最佳化器狀態低精度化的實驗。
- 集合通訊庫：NCCL、RCCL、Gloo 以及 SHARP 網路內計算，關係到 ZeRO 分片狀態的 AllGather 與 ReduceScatter 效率，是大規模最佳化器狀態通訊的硬體-軟體介面。
- 記憶體/儲存層級：HBM（高頻寬視訊記憶體）、CPU DRAM、NVMe SSD 的容量與頻寬，直接限制狀態解除安裝方案的可行性邊界。CXL 互聯記憶體等新興技術可能進一步改變上游硬體格局。
下游
最佳化器狀態管理與最佳化的下游覆蓋訓練棧的多個層次：
- 分散式訓練架構：DeepSpeed、Megatron-LM、Colossal-AI、PyTorch FSDP/FSDP2。它們將狀態分片、通訊、解除安裝編譯成易用介面，使普通開發者不必重寫最佳化器即可降低記憶體佔用。
- 量化訓練庫：bitsandbytes（8-bit Adam）、NVIDIA APEX、Transformer Engine 中的 FP8 最佳化器元件。它們提供量化的狀態儲存與反量化 kernel，在保持訓練精度的前提下大幅削減視訊記憶體。
- 雲端平台與 MLOps：AWS SageMaker、Google Cloud TPU VM、Azure ML 等平台將 DeepSpeed/FSDP 等整合到訓練作業模板中，並通過 API 暴露“最佳化器狀態解除安裝”等選項，影響訓練成本預估與作業排程。
- 學習率排程器與除錯工具：預熱（warmup）、餘弦退火等排程策略與最佳化器狀態 m、v 的動態耦合；TensorBoard、WandB 等工具允許記錄狀態範數、有效學習率等，幫助工程師除錯狀態健康度。
受益公司
以下公司及機構在最佳化器狀態相關技術棧中扮演重要角色，其行為影響產業的演進節奏（僅描述事實，不構成任何投資建議）。
- 微軟（DeepSpeed 團隊）：推出 ZeRO、ZeRO-Infinity、ZeRO++ 等一系列狀態分片與解除安裝方案，並積極將其整合至 Azure 機器學習服務和與 NVIDIA Megatron-LM 的協作中。
- Meta（FAIR & PyTorch 生態）：PyTorch 原生支援 Adam/AdamW，並通過 FSDP/FSDP2 將 ZeRO 理念深度融入 PyTorch 核心；2023 年釋出的 MeZO 零狀態微調技術由 Meta FAIR 提出，極具顛覆性。
- NVIDIA：提供 APEX 混合精度訓練工具、Transformer Engine FP8 支援，以及在 Megatron-LM 中實現高效的分散式最佳化器狀態管理；硬體方面，Hopper 架構通過 FP8 tensor core 和更大的 HBM 間接擴大狀態的可用容量與頻寬。
- Hugging Face：通過 Transformers、PEFT、Accelerate 等庫，將 bitsandbytes 8-bit 最佳化器無縫整合，大幅降低社群使用大型模型進行微調的門檻。
- Google：其 JAX 生態的 Optax 庫提供高度模組化的最佳化器狀態組合；TPU 訓練棧針對最佳化器狀態儲存和通訊有專門的架構最佳化，支撐了 PaLM 等大型模型的訓練。
- Colossal-AI（潞晨科技）：提供異構記憶體管理和最佳化器狀態分片方案，針對數百億引數模型的訓練進行工程簡化。
市場規模
最佳化器狀態並非一個可獨立採購的產品，它更像大型模型基建中的“隱性技術元件”。因此，公開資料未見針對最佳化器狀態技術獨立市場規模的統計。其商業影響力滲透在訓練硬體、雲端服務和開源生態之中：
- 視訊記憶體成本佔比：在典型的 100B+ 引數模型訓練中，最佳化器狀態佔用的視訊記憶體可達到總模型狀態（引數+梯度+狀態）的 50%–75%（據微軟 ZeRO 論文及 Meta 開源訓練日誌推算）。對於數千塊 GPU 的叢集，這部分視訊記憶體對應的硬體成本可高達數千萬美元。
- 門檻降低效果：ZeRO 系列技術已使訓練百億引數模型所需的最低 GPU 數量降至原來的 1/4 甚至更低（來源：DeepSpeed 公開 benchmark，2021—2023）。這意味著同等規模的實驗費用驟減，推動更多中小型公司參與大型模型創新。
- 8-bit 最佳化器使用量：bitsandbytes 的 8-bit Adam 被整合進 Hugging Face 生態後，截至 2023 年底，其通過 PyPI 和 Conda 的月下載量估計超過千萬次（綜合 GitHub star 增長曲線及 Hugging Face 社群帖文估算；精確資料未由單一機構揭露），已成為消費級顯示卡微調大型模型的標準配置。
- 雲端服務營收對映：主流雲端廠商的 GPU 訓練例項（如 AWS p4d/p5、Azure NDv5）銷量在 2023—2024 年快速增長，很大程度上由大型模型訓練推動，而最佳化器狀態記憶體最佳化是決定這些例項利用率的關鍵因素。不過，將雲端營收精確拆解至狀態最佳化技術上並不現實。
總體上，最佳化器狀態技術的市場價值體現在“省出來的 GPU 時費用”上，其間接經濟規模約為大型模型訓練總硬體/雲支出的重要組成部分，但缺乏獨立統計口徑。
玩家對比
在實際落地中，技術團隊通常會比較多種支援最佳化器狀態管理的架構或庫。以下從功能、效能、易用性等方面進行橫向對比（基於公開文件及社群反饋，截至 2024 年 Q2）。
| 玩家方案 | 狀態管理策略 | 支援的最大型模型規模 | 易用性 | 通訊與解除安裝特色 | 代表使用者 |
|---|---|---|---|---|---|
| DeepSpeed ZeRO (Stage 1-3) | 狀態/梯度/引數分片 + CPU/NVMe 解除安裝 | 實測 1000B+ 引數（ZeRO-Infinity） | 通過 json 配置檔案與單行 API 呼叫，高度整合 | 高度最佳化的 AllGather/ReduceScatter；支援與 Megatron 混合並行；ZeRO-Infinity 支援 NVMe 解除安裝 | GPT-NeoX、BLOOM、多個開源百億模型 |
| PyTorch FSDP / FSDP2 | 引數、梯度和狀態的統一分片（類似 ZeRO-3） | 數百 B 引數（Facebook 內部） | 原生 PyTorch 包裝器，config 簡單，社群支援活躍 | 通訊與計算重疊優秀；FSDP2 支援按 Transformer 層分片，減少通訊空隙 | Llama 2/3 訓練、Vicuna 等 |
| Colossal-AI | 異構記憶體（GPU-CPU-NVMe）統一管理；支援 ZeRO 風格分片及創新的 Gemini 最佳化 | 達到百 B 級別 | 提供配置檔案和少量程式碼改動 | 細粒度的異構記憶體排程，可自定義狀態放置策略 | 某國產開源大型模型 |
| bitsandbytes 8-bit Adam | 8-bit 塊量化儲存 + fp32 反量化計算 | 僅影響單卡，最佳化器狀態降至 0.5×，使單卡可微調 70B 模型 | 一行程式碼替換最佳化器類，與 Hugging Face Trainer 無縫銜接 | 無分散式分片，純單卡方案 | 社群微調 Llama、Falcon 等 |
| MeZO (Meta FAIR) | 零狀態，基於 SPSA 隨機梯度估計 | 單卡可微調 OPT-66B 等，記憶體節約 12× 以上（相對於 Adam） | 需要修改訓練迴圈，目前為學術程式碼 | 無需反向傳播，僅需兩次前向，記憶體足跡極小 | 研究機構、邊緣場景探索 |
| NVIDIA Transformer Engine | FP8 梯度與狀態（實驗性） | 與 Megatron-LM 結合的百億級訓練 | 依賴於 NVIDIA NGC 容器，對硬體有要求 | 利用 FP8 實現 2× 頻寬提升，狀態壓縮尚在完善中 | 部分 NVIDIA 合作伙伴 |
選擇的關鍵權衡在於：訓練規模 vs. 吞吐效率 vs. 精度的損失容忍度。生產環境中通常採用 DeepSpeed ZeRO 或 FSDP 結合 8-bit 最佳化器，以求在擴充套件性和訓練速度間取得平衡。
風險
最佳化器狀態相關技術的部署和演進面臨以下風險：
- 狀態記憶體牆持續存在：模型引數量每代增長 10×，而 GPU 視訊記憶體增長僅 2× 左右（據 NVIDIA 歷代產品引數），即便有分片和解除安裝，超大型模型訓練仍需巨大的叢集規模，抬升硬體投資和故障恢復成本。
- 量化狀態導致的精度塌陷：8-bit 最佳化器在部分任務中可能出現收斂變慢、最終精度小幅下降或訓練不穩定（尤其在強化學習、GAN 等對梯度質量敏感的場景中）。FP8 狀態尚處實驗階段，廣泛用於最佳化器狀態存在未知的長期隱患。
- 分散式通訊開銷：ZeRO 分片雖降低單卡記憶體，但引入額外的 AllGather/ReduceScatter 通訊。隨著模型增大和卡數增多，通訊時間可能超過計算，成為瓶頸。若網路拓撲不佳，狀態分片的收益會被通訊延遲侵蝕。
- 零狀態方案收斂限制：MeZO 等基於隨機擾動的方法，梯度估計方差遠高於精確反向傳播，導致達到同等精度所需訓練步數顯著增加（論文顯示可能 10—100 倍），只適用於微調場景，完全不適合從頭預訓練。
- 軟硬體生態鎖定：高階狀態最佳化功能（如 ZeRO-Infinity 解除安裝、FP8 狀態）高度耦合於特定架構和硬體（如 NVIDIA GPU），可能削弱模型的可移植性，增加供應鏈風險。
- 誤用風險：缺乏理解的狀態配置（如不當的 β₁、β₂、ZeRO 階段選擇）可能導致隱性發散或資源浪費，而除錯最佳化器狀態需要較強的專家知識，人才稀缺。
誤讀糾偏
誤讀 1：“Adam 的 m 和 v 就是多步梯度的累加，所以最佳化器狀態等同於梯度累積。” 糾正：m 是梯度的一階矩估計，類似帶摩擦的“速度”；v 是梯度的二階矩估計，衡量近期梯度波動的“熱度”。梯度累積則是顯式地對多步梯度求均值，並不產生持久狀態，只是延遲最佳化器步驟。兩者機制完全不同。
誤讀 2：“既然有分片和解除安裝，最佳化器狀態就不再是瓶頸，可以無限制增大型模型。” 糾正：分片和解除安裝只是將狀態分散儲存，但單卡仍需儲存部分引數、梯度及啟用等；狀態通訊會佔用互聯頻寬，在超大規模下通訊瓶頸依舊存在。物理上，引數和啟用的儲存需求仍會隨模型規模線性增長，不能無限擴充套件。
誤讀 3：“零狀態方案（MeZO）可以完全替代 Adam，為大型模型訓練省下所有狀態記憶體。” 糾正：MeZO 免去了狀態記憶體，但其梯度估計方差大，收斂所需的迭代步數遠多於 Adam，總計算量可能反而更高。目前實證表明它適合於記憶體極度受限的微調場景，距離從零開始的預訓練仍有顯著差距。
誤讀 4：“8-bit 最佳化器只是把狀態存為 8-bit，計算時還是 fp32，所以速度會大幅下降。” 糾正：8-bit 最佳化器在更新時確實需要反量化到 fp32 進行計算，然後重新量化，這引入了額外 kernel 開銷。但工程實現高度最佳化，整體訓練吞吐量下降幅度通常