⏶189
GaLore:通过梯度低秩投影实现内存高效的LLM训练
03月06日发表
04月12日由
AK 提交

作者:
Jiawei Zhao,
Zhenyu Zhang,
Beidi Chen, Zhangyang Wang,
Anima Anandkumar,
Yuandong Tian

摘要
训练大型语言模型 (LLM) 带来了巨大的内存挑战,这主要是由于权重和优化器状态的规模不断增长。常见的内存减少方法(例如低秩自适应 (LoRA))在每层中向冻结的预训练权重添加可训练的低秩矩阵,从而减少可训练参数和优化器状态。然而,这种方法通常在预训练和微调阶段都表现不佳,因为它们将参数搜索限制在低秩子空间并改变了训练动态,并且进一步地,可能需要全秩热启动。在这项工作中,我们提出了梯度低秩投影 (GaLore),这是一种训练策略,允许全参数学习,但比常见的低秩自适应方法(如 LoRA)更节省内存。我们的方法在优化器状态中最多可减少 65.5% 的内存使用量,同时保持了在 LLaMA 1B 和 7B 架构上使用 C4 数据集进行预训练(最多 19.7B 个 token)以及在 GLUE 任务上微调 RoBERTa 的效率和性能。我们的 8 位 GaLore 进一步将优化器内存减少高达 82.5%,并将总训练内存减少 63.3%,与 BF16 基线相比。值得注意的是,我们首次证明了在具有 24GB 内存的消费级 GPU(例如,NVIDIA RTX 4090)上预训练 7B 模型的可行性,而无需模型并行、检查点或卸载策略。
评论
对于图 1。
您能否也包括 8bit-Adam + 每层权重更新但在梯度更新中不进行秩缩减会对内存使用产生什么影响?根据 Lomo 论文 / https://arxiv.org/abs/2306.09782,似乎它也会显着减少内存使用的浅绿色部分,因为梯度在每层立即被消耗+丢弃?
论文作者
感谢您的评论!我们在这里有第三方评估:https://github.com/jiaweizzhao/GaLore/issues/6。仅 GaLore(不进行每层权重更新)就具有与每层权重更新相当的内存减少量。它们是正交技术。通过将它们组合在一起,您可以在 24G 内存(例如 4090)内运行 7B 预训练。
GaLore:使用内存高效的梯度投影彻底改变 LLM 训练
链接 🔗:
👉 订阅: https://www.youtube.com/@Arxflix
👉 Twitter: https://x.com/arxflix
👉 LMNT (合作伙伴): https://lmnt.com/
作者:Arxflix
希望澄清图 1:这里的批量大小、序列长度和词汇表大小是多少?因为我预计激活会占用更多空间...
根据图 1 的标题,批量大小似乎为 256
根据脚注 1,序列长度似乎为 2048
根据 config from repo,词汇表大小为 32000
根据脚注 2,使用 bf16,因此每个浮点数为 2 个字节
因此,只有模型的 logits 应该占用 256 * 2048 * 32000 * 2
字节或 31.25 GB。图 1 中哪里需要这个内存?
谢谢!
https://unrollnow.com/status/1765699202489651364 用于评论/TL;DR 线程(或原始线程:https://twitter.com/BlackHC/status/17656992024896513640)