「田渊栋领衔」GPU工业要变天了:GaLore,用3/1 的显存,预训练大语言模型



本文为1027字,建议阅读3分钟
这下显卡工业又要变天了? GaLore,用3/1 的显存,预训练大语言模型, 将权重更新的gradient投影到低维空间,减少矩阵大小,减少内存消耗, 感觉是一个很数学的工作
加州理工学院、Meta AI(田渊栋)、德克萨斯大学奥斯汀分校和卡内基梅隆大学的研究人员共同发表论文GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection(GaLore:通过梯度低秩投影实现内存高效的大型语言模型训练)

研究人员提出GaLore技术引起了业界巨大关注,它通过梯度低秩投影(Gradient Low-Rank Projection)的方法实现对大型语言模型训练过程中的内存使用量大幅度削减,同时保持了预训练和微调任务的性能与效率
在训练大型语言模型(LLMs)时,我们面临着内存使用的越来越严峻挑战,这主要是因为模型权重和优化器状态的规模不断扩大。传统的内存优化方法,比如低秩适应(LoRA),通过在每层的预训练权重中加入一个可训练的低秩矩阵,来减少可训练参数和优化器状态的数量。但这些方法在预训练和微调阶段往往无法达到全秩权重训练的性能,因为它们将参数搜索限制在了一个低秩子空间,并改变了训练过程的动态,而且可能还需要从全秩状态开始训练。
GaLore 采用了一种与 LoRA 和 ReLoRA 等传统方法截然不同的新策略。它不是在权重空间进行低秩投影(W = W_0 + A@B),并且每隔 T 步将这个投影整合到权重中,而是在梯度空间执行低秩投影:G = P @ G' @ Q^T。这种方法在处理梯度时更为高效,有助于减少内存的使用,同时保持模型训练的质量和速度
在这项研究中,采用GaLore的训练策略,它允许我们进行全参数学习,这种方法在优化器状态上最多可以减少65.5%的内存使用,同时在预训练LLaMA 1B和7B架构以及使用C4数据集处理高达19.7亿个令牌时,以及在微调RoBERTa处理GLUE任务时,都能保持高效和性能。8位GaLore版本进一步将优化器内存减少了高达82.5%,总训练内存减少了63.3%,与BF16基线相比。

特别值得一提的是,GaLore方法首次证明了在没有采用模型并行、检查点保存或数据卸载策略的情况下,可以在拥有24GB内存的消费级GPU(如NVIDIA RTX 4090)上预训练一个7B参数的模型
结语
通过GaLore,我们可以在内存较低的硬件设备上训练更大的模型。未来的研究方向包括但不限于将GaLore应用于其他类型的模型(如视觉变压器和扩散模型)训练,探寻利用量化或特殊参数化手段实现更低内存投影矩阵的方法,以及在消费级低带宽硬件上实现弹性分布式训练的可能性。随着GaLore这类工具的出现和发展,科研社区和个人有望以更少的资源在消费级硬件上高效地训练出更大规模的语言模型。
⭐星标AI寒武纪,好内容不错过⭐
用你的赞和在看告诉我~

单个GPU搓大模型?👇👇
到顶部