高端显卡也能训练模型，AI研究人员提出“GaLore”减少内存...

darthracer · 2024-4-15 23:52

随着大型语言模型的规模越来越大，若对所有模型参数进行训练，成本势必非常高昂，研究人员设法减少内存使用量，其中新方法GaLore（Gradient Low-Rank Projection）可让训练过程更省内存，甚至用户中计算机搭配高端显卡就能训练。

训练大型语言模型为运算基础设施带来无法忽视的内存挑战，主要是因权重大小和优化器（optimizer）状态不断增加。常见内存减少方法，例如微软研究团队所提出LoRA（Low-Rank Adaptation），概念是冻结原本预训练模型权重，减少训练参数量和优化器状态，也能达到很好的微调（fine-tuning）效果。

不过，LoRA通常在预训练和微调阶段表现不佳，因为它将参数搜索限制在low-rank subspace并改变训练动态，性能通常比原本训练方法差。

日前发布的《GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection》新论文，研究人员提出GaLore是一种新的模型训练策略，可让模型使用全部参数进行学习，同时比LoRA更省内存。

GaLore在优化器状态下将内存使用量减少多达65.5%，同时在LLaMA 1B和7B架构下使用最多19.7B token的C4数据组进行预训练，依然保持良好性能，及在GLUE任务上微调RoBERTa的效率和性能。与BF16基准相比，8位元的GaLore进一步减少优化器内存多达82.5%，总训练内存则减少63.3%。

GaLore还有一大优点，是让原本只能用在大型服务器上的大型语言模型，通过家用计算机的高端显卡（例如NVIDIA RTX 4090）也能进行训练，研究人员首次证明在具有24GB内存的GPU预训练7B参数模型是可行的。