[控场AI]
· 3 分钟阅读· 1,938 字

ResQAdam:残差保留的自适应优化器,大幅降低QAT内存开销

ResQAdam:残差保留的自适应优化器,大幅降低QAT内存开销

ResQAdam 用单个 BF16 动量与残差符号更新,将 QAT 优化器内存降低约七成且精度接近 AdamW。

量化感知训练(QAT)是获得高精度低比特大语言模型最可靠的方法,但 AdamW 需为每个权重维护两个 FP32 动量状态,内存开销巨大。低秩优化器(如 GaLore、LDAdam)虽能压缩状态,却在 QAT 中精度明显下降,原因是 QAT 动量有相当部分落在 top-r 奇异子空间之外,被投影方法逐步丢弃。论文发现即使只保留这部分残差的符号也仍然有用,据此提出 ResQAdam:仅维护一个 BF16 动量,对主子空间的 r×r 坐标做 Adam 更新,并叠加由残差计算出的归一化符号更新。在 2 比特量化、两种从头 QAT 方法与三种微调流程上的实验显示,优化器状态内存相对 AdamW 减少 71%–74%,多数设置下精度接近 AdamW,吞吐量下降不到 2%。

量化感知训练(QAT)是获得高精度低比特大语言模型(LLM)最可靠的方法,但使用AdamW进行训练时,每个权重矩阵都需要维护两个密集的FP32状态,带来巨大的内存开销。近日,一篇题为《ResQAdam: Residual-Preserving Adaptive Optimization for Memory-Efficient Quantization Aware Training》的论文提出了一种新优化器ResQAdam,旨在解决这一痛点。

ResQAdam论文标题

QAT的内存瓶颈与低秩优化器的局限

QAT通过在训练过程中模拟量化误差,让模型适应低精度表示。然而,AdamW优化器需要为每个参数存储一阶和二阶动量,均为FP32精度,这使得训练内存需求激增。低秩内存高效优化器(如GaLore和LDAdam)通过将梯度投影到低秩子空间来减少状态存储,但在QAT场景下表现显著下降。

论文作者发现,QAT动量中有相当一部分位于其top-r奇异子空间之外,而基于投影的方法在每一步都会丢弃这部分信息。更关键的是,他们进一步观察到,即使只保留残差的符号,这部分信息仍然有用。这一发现为设计新优化器提供了思路。

奇异值分解与主奇异子空间:对梯度(或动量)矩阵做奇异值分解(SVD)后,会得到一组按奇异值从大到小排列的方向,其中前 r 个最大的奇异值对应的方向张成的空间称为主(top-r)奇异子空间。低秩优化器的核心假设是:优化过程中最重要的信息集中在这个子空间里,因此只需在该子空间中存储和更新状态即可大幅节省内存。然而在 QAT 场景下,量化误差带来的梯度扰动会使一部分有用信息分布在主奇异子空间之外,即所谓“残差”部分;投影类方法在每一步都会截断这部分信息,累积起来便造成精度损失。ResQAdam 的关键观察正是:残差虽不能全部存储,但其符号方向仍携带有效信息,保留符号即可在几乎不增加内存的前提下挽回大部分损失。

BF16 与 FP32 的精度差异:BF16(bfloat16)与 FP32 同为 16 位与 32 位浮点格式,但 BF16 只有 8 位尾数(FP32 有 23 位),动态范围与 FP32 相同而精度低得多。优化器状态从 FP32 改为 BF16 可直接将这部分内存减半,是内存高效训练中常用的手段;但 BF16 动量在累积微小更新时容易丢失低位信息,因此通常需要配合归一化、符号化等补偿策略,这也是 ResQAdam 仅对主子空间坐标做 Adam 更新、对残差只保留符号的原因之一。

ResQAdam的设计:保留残差符号的自适应更新

基于上述观察,研究者提出了ResQAdam。它只维护一个BF16精度的动量,仅对主奇异子空间的r×r坐标应用Adam更新,并添加一个从残差计算出的归一化符号更新。具体来说,ResQAdam将动量分解为主子空间和残差两部分:对主子空间部分使用标准的Adam自适应学习率,而对残差部分则仅保留其符号方向,进行归一化后加到更新中。这样既保留了残差中的有效信息,又避免了存储完整的动量状态。

Adam 的自适应学习率与归一化符号更新:Adam 为每个参数维护一阶动量(梯度均值)和二阶动量(梯度平方均值),更新量近似为 m/√v,即用二阶动量对一阶动量做归一化,使不同量级的参数获得大致相当的步长。ResQAdam 对主子空间的 r×r 坐标沿用这一机制;而对残差部分,由于无法承担存储完整二阶动量的开销,改为只取残差的方向(符号),再做归一化后与主子空间更新相加。这相当于用“signSGD 式”的符号更新近似 Adam 的归一化效果——符号更新本身具有对梯度尺度不敏感的特性,与 Adam 的尺度归一化在直觉上一致,因此能在几乎没有额外状态的情况下保留残差中的有效信息。

实验评估:内存节省显著,精度接近AdamW

作者使用两种从头开始的QAT方法和三种QAT微调流程进行评估,并将权重量化到2比特。结果显示,ResQAdam将优化器状态内存相对于AdamW减少了71%至74%,同时在大多数设置下精度接近AdamW,训练吞吐量下降不到2%。这意味着ResQAdam在显著降低内存需求的同时,几乎不影响训练速度和模型质量。

总结与展望

ResQAdam通过保留残差符号信息的创新设计,在低秩优化框架下实现了接近AdamW的QAT性能,大幅降低了内存占用。这对于在有限硬件资源下训练低比特大语言模型很关键。未来,该方法有望进一步推广到其他内存敏感的训练场景中。

#ResQAdam#量化感知训练#内存高效优化器#低秩优化#QAT#大语言模型#AdamW#低比特量化
分享:

相关推荐