优化器状态
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 开销。但工程实现高度优化,整体训练吞吐量下降幅度通常