提出新优化器,让大模型在内存紧张时仍能高效训练。
Backward-Friendly Optimization: Training Large Language Models with Approximate Gradients under Memory Constraints
- 用低秩雅可比近似压缩梯度信号,减少内存占用。
- 通过误差反馈补偿近似误差,保证训练收敛性。
- 无需改模型结构,适合资源受限的LLM微调场景。
大型语言模型(LLMs)全量微调通常内存开销巨大,因传统优化器如SGD或Adam依赖缓存的中间激活值来计算精确梯度。现有方法或修改模型架构(如可逆网络),或以增加计算代价换取内存节省(如激活检查点),但优化器本身未被改进。本文提出GradLite,一种面向反向传播的新型优化器,放宽对精确梯度的要求,即使中间激活值被大幅丢弃或近似,也能实现高效训练。GradLite采用两项关键技术:(i) 低秩雅可比近似,降低反向传播误差信号维度;(ii) 误差反馈校正,累积并补偿迭代中的近似误差,确保梯度估计无偏且方差有界,从而保持与Adam相当的收敛速度。理论分析表明,该方法可保障收敛性。实验显示,相比检查点法和优化器中心型基线(LoMo、GaLore),GradLite将优化器状态与激活内存消耗降低最高50%,在推理(MMLU、GSM8K)、多语言及对话任务上性能相当或更优。
原文摘要 · Abstract (English)
Full fine-tuning of Large Language Models (LLMs) is notoriously memory-intensive, primarily because conventional optimizers such as SGD or Adam assume access to exact gradients derived from cached activations. Existing solutions either alter the model architecture (e.g., reversible networks) or trade memory for computation (e.g., activation checkpointing), but the optimizer itself remains untouched. In this work, we introduce GradLite, a backward-friendly optimizer that relaxes the requirement of exact gradients, enabling efficient training even when intermediate activations are aggressively discarded or approximated. GradLite leverages two key techniques: (i) low-rank Jacobian approximation, which reduces the dimensionality of backpropagated error signals, and (ii) error-feedback correction, which accumulates and compensates approximation errors across iterations to preserve convergence guarantees. We provide a theoretical analysis showing that GradLite maintains unbiased gradient estimates with bounded variance, ensuring convergence rates comparable to Adam. Empirically, GradLite reduces optimizer-state and activation memory consumption by up to 50\% without architectural changes, and achieves on-par or superior downstream performance on reasoning (MMLU, GSM8K), multilingual, and dialogue benchmarks compared to checkpointing and optimizer-centric baselines (LoMo, GaLore).
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。