用小波变换压缩大模型训练状态,省内存还提速。
GWT: Scalable Optimizer State Compression for Large Language Model Training
- 用无SVD的小波变换压缩梯度矩阵,保留关键信息。
- 内存减少52%,在LLaMA-3B上速度提升1.9倍。
- 兼容Adam等多类优化器,适合大规模模型训练场景。
训练大语言模型需大量内存,其中自适应优化器(如Adam)维护的动量状态占主要部分。现有方法通过低秩投影压缩参数或梯度,可能丢失子空间外信息,或因子空间构建和奇异值分解(SVD)增加开销。本文提出梯度小波变换(GWT),一种无需SVD的框架,对矩阵梯度应用多级Haar变换。GWT仅存储压缩系数的均值与方差状态,暂存所有细节系数。理论证明该正交变换保持完整当前梯度,并在有界预条件和可直接验证的粗动量相干条件下,保证梯度对齐与平稳点收敛性。在语言模型预训练与下游微调实验中,GWT将模型及优化器状态内存减少最多52%,同时保持或提升模型质量。在LLaMA-3B预训练中,其吞吐量达8位Adam的1.9倍,且兼容Adam、Adam-mini、Muon等多种优化器。
原文摘要 · Abstract (English)
Training large language models (LLMs) requires substantial memory, a significant fraction of which is consumed by the moment states maintained by adaptive optimizers such as Adam. Existing memory-efficient approaches commonly compress parameters or gradients through low-rank projections, which may discard information outside the selected subspace or incur additional costs from subspace construction and singular value decomposition (SVD). We introduce \textbf{Gradient Wavelet Transform (GWT)}, an SVD-free framework that applies a multilevel Haar transform to matrix gradients. GWT maintains first- and second-moment states only for the compact approximation coefficients, while retaining all detail coefficients transiently in each update. We show that the orthogonal transform preserves the complete current gradient and establish gradient alignment and a stationary-point guarantee for smooth objectives under bounded preconditioning and a directly verifiable coarse-momentum coherence condition. Experiments on language-model pre-training and downstream fine-tuning show that GWT reduces estimated model-and-optimizer-state memory by up to 52\% while maintaining competitive or improved model quality. On LLaMA-3B pre-training, GWT achieves $1.9\times$ the throughput of 8-bit Adam and is compatible with multiple optimizer families, including Adam, Adam-mini, and Muon.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。