ResQAdam:低显存量化感知训练优化器, optimizer 状态省 71%
ResQAdam: Residual-Preserving Adaptive Optimization for Memory-Efficient Quantization Aware Training
研究团队发现 GaLore 这类低秩优化器做 2 bit 量化训练会掉点,于是提出 ResQAdam,优化器显存省七成多,速度几乎不掉。
ResQAdam 是一种面向量化感知训练(QAT)的低显存优化器。GaLore 和 LDAdam 这类低秩优化器在 QAT 中会因丢弃主子空间外的动量而明显退化,ResQAdam 保留单个 BF16 动量,对主子空间的 r×r 坐标应用 Adam,并对残差添加归一化符号更新。在权重低至 2 bit 的两种从头 QAT 方法和三种 QAT 微调流水线中,ResQAdam 相比 AdamW 减少 71% 到 74% 的优化器状态显存,训练吞吐量下降不到 2%。
ResQAdam: Residual-Preserving Adaptive Optimization for Memory-Efficient Quantization Aware Training
Quantization-aware training (QAT) is the most reliable way to obtain accurate low-bit large language models (LLMs), but training with AdamW keeps two dense FP32 states for every weight matrix. We find that low-rank memory-efficient optimizers such as GaLore and LDAdam degrade substantially in QAT. A considerable part of the QAT momentum lies outside its top-r singular subspace, and projection-based methods discard this part at every step. We further observe that this residual remains useful when only its sign is kept. Based on these observations, we propose ResQAdam. It keeps a single BF16 momentum, applies Adam only to the r x r coordinates of the principal subspace, and adds a normalized sign update computed from the residual. We evaluate ResQAdam with two from-scratch QAT methods and three QAT fine-tuning pipelines, with weights quantized down to 2 bits. It reduces optimizer-state memory by 71% to 74% relative to AdamW, stays close to AdamW in most settings, and lowers training throughput by less than 2%.