动态调整训练中优化器精度,大幅降低显存占用且不损失模型效果。
STQuant: Spatio-Temporal Adaptive Framework for Optimizer Quantization in Large Multimodal Model Training

- 按层、变量和训练阶段动态分配量化精度,避免固定精度的误差积累。
- 在GPT-2和ViT上实现84.4%的显存压缩,平均仅需5.1位。
- 算法高效稳定,计算开销仅为O(N/K),额外内存需求为O(1)。
量化是降低大规模模型训练内存开销的有效方法。然而,现有方法多采用固定精度策略,忽视了优化器状态分布随层和训练步骤显著变化的事实,导致明显精度下降。为突破固定量化限制,我们提出STQuant,一个分布式训练框架,通过跨层、状态变量和训练步骤的动态精度分配,降低优化器状态内存占用,同时保持模型质量。直接应用动态量化面临两大挑战:一是优化器状态数值敏感,量化噪声易引发训练不稳定;二是多状态与多层组合带来巨大搜索空间。STQuant通过两项关键技术解决:1)一种可证明近优的因子选择策略,精准识别影响精度适配的关键因素;2)一种动态转换决策算法,将搜索复杂度从指数级降至线性。在GPT-2和ViT上的实验表明,STQuant将优化器状态内存减少84.4%,平均位宽低至5.1比特,优于现有方案。此外,其计算开销仅为O(N/K),额外内存需求为O(1)。
原文摘要 · Abstract (English)
Quantization is an effective way to reduce the memory cost of large-scale model training. However, most existing methods adopt fixed-precision policies, which ignore the fact that optimizer-state distributions vary significantly across layers and training steps. Such uniform designs often introduce noticeable accuracy degradation. To move beyond fixed quantization, we propose STQuant, a distributed training framework that reduces the memory footprint of optimizer states via dynamic precision allocation across layers, state variables, and training steps, while maintaining model quality. Naively applying dynamic quantization during training is challenging for two reasons. First, optimizer states are numerically sensitive, and quantization noise can destabilize quality. Second, jointly considering multiple states and layers induces a large combinatorial search space. STQuant addresses these challenges with two key techniques: 1) a provably near-optimal factor selection strategy that accurately identifies the most influential factors for precision adaptation. 2) a dynamic transition decision algorithm that reduces the search cost from exponential to linear complexity. Experiments on GPT-2 and ViT show that STQuant reduces optimizer-state memory by 84.4%, achieving an average bit-width of as low as 5.1 bits, compared with existing solutions. Moreover, STQuant incurs only O(N/K) computational overhead and requires O(1) extra space.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。