女科学家提出GaLore：为消费级GPU上高效训练LLM铺平道路

Anima Anandkumar，这位研究科学家堪称人工智能界的传奇人物。作为加州理工学院的布伦教授(Bren Professor)，她是ACM、IEEE、Schmidt Science、Guggenheim以及Alfred.P.Sloan基金会的成员，堪称学术界的一颗耀眼明星。她曾获得NeurIPS最佳论文奖和ACM Gordon Bell特别奖，展现出在科学研究领域的卓越才华。

在职业生涯的早期，Anima曾担任过AWS首席科学家，为亚马逊的科技发展贡献了巨大的力量。后来她加入了NVIDIA，担任人工智能研究的高级总监，整整5年3个月，为该公司的技术创新和发展贡献了自己的智慧。

然而，就在去年11月，当NVIDIA还在被巨大光环所笼罩的时刻，Anima选择离开NVIDIA，重新回到学术界，回归加州理工学院的教学和研究工作。

她的职业经历不仅丰富多彩，更是为人工智能领域树立了一个标杆。她的故事告诉我们，才华与毅力能够让我们在不同领域中取得成功，不断追求卓越，永不止步。

就在3月8日这一天，大家都在通过各种活动展现女性在各行各业做出的贡献之时，Anima在社交平台上公布了一个贴文，迅速轰动国外AI圈：

我们首次证明，LLaMA 7B可以在仅具有24GB内存的单一消费级GPU(RTX 4090)上进行预训练(不是微调!)。这意味着训练期间用于存储优化器状态的内存减少了82.5%以上。

在帖文里，Anima解释说：在单张消费级N卡上预训练拥有70亿参数的大型语言模型(LLM)可能是一个令人惊讶的概念。预训练和微调LLM不仅需要大量算力，还需要大量的内存来存储运算数据，包括数十亿个可训练参数、它们的梯度和优化器状态。例如，使用单一批次大小从头开始预训练LLaMA 7B模型需要至少58GB内存(14GB用于可训练参数，42GB用于Adam优化器状态和权重梯度，2GB用于启动)，这使得此任务在消费级GPU(例如具有24GB内存的NVIDIA RTX 4090)上不可行。

为了解决GPU内存不足的困境，陆续出现了LoRA这类的训练参数缩减方法，将可训练的低秩矩阵添加到每层冻结的预训练权重中，从而减少训练参数和优化器状态。然而，LoRA只能用在微调，无法用于全参数的LLM预训练。为了在内存受限的硬件上实现LLM预训练任务，Anima与她的Caltech同事提出了全参数训练策略——GaLore(Gradient Low-Rank Projection)。

在这项工作中，Anima的团队提出了梯度低秩投影(GaLore)，这是一种训练策略，可以进行全参数学习，但比常见的低秩适应方法(如LoRA)更节省内存。他们的方法在优化器状态中将内存使用量降低了高达65.5%，同时在使用C4数据集进行LLaMA 1B和7B架构的预训练以及在GLUE任务上对RoBERTa进行微调时，保持了效率和性能。

此外，团队的8位GaLore进一步将优化器内存减少了高达82.5%，与BF16基准相比，将总训练内存减少了63.3%。值得注意的是，他们首次证明，在消费级GPU上(例如NVIDIA RTX 4090)，不需要模型并行、检查点或卸载策略，即可对7B模型进行预训练的可行性。

图：在单个设备上预训练LLaMA 7B模型，令牌批量大小为256，不使用激活检查点和内存卸载

彷佛给我们这些GPU卡穷人带来了一丝曙光。

在论文里，团队确定了GaLore的几个开放性问题，包括(1)将GaLore应用于其他类型模型的训练，如视觉变换器和扩散模型，(2)通过使用低内存投影矩阵(通过量化或特殊参数化)进一步提高内存效率，以及(3)探索在低带宽消费级硬件上进行弹性数据分布式训练的可能性。

最终，Anima和她的团队希望这项工作能够激发未来从低秩梯度投影的角度出发，对内存高效的LLM训练策略进行研究。她和他们相信，GaLore将成为社区在使用消费级硬件和有限资源训练大型语言模型的宝贵工具。

论文地址：https://arxiv.org/pdf/2403.03507.pdf

Github地址：

https://github.com/jiaweizzhao/galore?fbclid=IwAR3gQ7PDX45r4wqFVobfL6FTeTPGA5ohpBDgt7jwZTABaNF5KY19ttCsPoQ