彻底疯狂！实操用Golare在单块4090上预训练大模型及实测效率

在单块4090上预训练大模型？？

田渊栋等人新作：突破内存瓶颈，GaLore实现22G显存预训练大模型

最近正想着微调scGPT这种生物大模型，把那篇文章的小部分结果复现一下，结果24G显存还没推理就帕的一下装满了，看着手里的双4090只能本地部署部署的14B，微调微调还要并行运行，可怜的不要不要的，早知道多整几块2080Ti魔改了。

詹姆斯表示2080 22g是一张极为......

正当我纳闷的时候，想着有没有刚好4090就能跑下的模型，不浪费一点内存。还真不是开玩笑，大佬Meta FAIR 田渊栋的论文《 MobileLLM: Optimizing Sub-billion Parameter Language Models for On-Device Use Cases》就提供了GaLore架构，不牺牲模型性能的前提下显著降低内存消耗，主打在移动设备上运行10亿以下参数LLM。最重要的是，它能实现实现22G显存预训练大模型。这种送上门来的好项目怎么能不去骗去偷袭呢？

文章的预印本：https://arxiv.org/pdf/2402.14905.pdf

论文背景

LLM（Large Language Models）在运行时对内存的要求极高，不仅涵盖了模型参数本身的海量存储需求（常常以亿计），还包括了诸如Adam优化器所维护的梯度动量和方差等额外状态信息。例如，在初始预训练阶段，LLaMA 7B模型每批次仅处理一个样本时，其内存消耗就相当可观：大约58GB的总内存使用量中，有14GB专门用于存储模型可训练参数，而高达42GB的空间被Adam优化器的状态变量及权重梯度占据，另外还有2GB用于临时的激活函数输出。这意味着，对于拥有像NVIDIA RTX 4090这种仅配备24GB显存的消费级GPU而言，直接进行LLM的训练显然是不够的，因其内存资源远不足以容纳如此庞大的计算需求。

而通过梯度低秩投影，GaLore可以在不牺牲模型性能的前提下显著降低内存消耗，使得在消费级GPU上训练大型语言模型成为可能。田大佬本人强调：我们没有像 LoRA 那样假设低秩权重结构，而是证明权重梯度自然是低秩的，因此可以投影到（变化的）低维空间中。因此，我们同时节省了梯度、Adam 动量和方差的内存。 那么，到底是如何让显存大大降低呢？

算法策略

本篇论文核心提出了一种名为“GaLore”的梯度低秩投影训练策略，用于提高大型语言模型（LLMs）的内存效率。算法的基本思想是利用权重矩阵梯度在训练过程中表现出的缓慢变化的低秩结构，而不是直接将权重矩阵近似为低秩矩阵。具体做法如下：

首先，理论分析表明权重矩阵的梯度G具有低秩特征。基于这一发现，GaLore计算两个投影矩阵P∈Rm×r和Q∈Rn×r，用以将梯度矩阵G投影到低秩形式P⊤GQ，从而大幅减少依赖于梯度分量统计信息的优化器状态所占用的内存

对于链式可逆神经网络，其权重矩阵Wl在第l层的梯度Gl可以用特定形式表示定理3.2，其中包含前向传播中的雅可比矩阵及其乘积，以及与输出相关的项。

定理3.2（可逆模型的梯度形式）： 在具有ℓ2损失函数φ:= 1/2 ||y - N(x)||^2的链式可逆神经网络N(x) := NL(NL-1(...N1(x)))中， 在批量大小为1的情况下，第l层权重矩阵Wl的梯度Gl有以下形式：

其中：

 

请注意，这里假设了在softmax目标函数下梯度也有类似的结构，并适用于类似的情况。

https://github.com/jiaweizzhao/GaLore

GaLore的核心操作是在每个预设的频率T时刻重新计算投影矩阵P和Q，通过在不同的低秩子空间中进行更新来逐步逼近最优解。 例如，在时间步t1∈[0, T1-1]期间，权重W在由固定投影矩阵Pt1和Qt1确定的低秩子空间中根据投影梯度G˜t1进行更新； 当达到时间步T1后，会重新计算Pt2和Qt2进入下一个低秩子空间，并重复此过程直至收敛。

同时作者给出了梯度稳定秩随时间变化的不等式以及梯度随时间指数衰减至低秩的证明，这有助于理解为何GaLore能够在保证性能的同时减少内存需求。在附带的Adam优化器结合GaLore的算法描述中，即每经过一定频率T的时间步长，根据当前梯度Gt重新初始化投影矩阵，并基于此进行低秩投影后的梯度更新，以实现高效、全参数学习下的内存节省。

Adam with GaLore 结合了Adam优化器和梯度低秩投影技术，下面是结合这两种方法进行训练的简化算法流程：

• 首先，按照标准Adam更新规则初始化一阶动量M0和二阶动量V0，然后开始迭代训练过程：

• 获取当前时间步t的权重矩阵Wt对应的梯度Gt = -∇Wφt(Wt)。

• 如果t满足一定的子空间频率条件（例如t mod T= 0），则执行以下操作：

• 对梯度矩阵Gt进行奇异值分解（SVD），得到U, S, V三个矩阵。

• 初始化左投影矩阵Pt为U[:, :r]（若m <= n，否则选择合适的投影矩阵）。

• 否则，复用上一步的投影矩阵Pt = Pt-1。

• 将梯度投影到紧凑空间得到Rt = P⊤ t Gt。

• 使用Adam或其他相似优化器更新投影后的梯度Rt：

• 更新一阶动量Mt：Mt ← β1·Mt-1 + (1 - β1)·Rt

• 更新二阶动量Vt：Vt ← β2·Vt-1 + (1 - β2)·Rt²

• 计算归一化的低秩梯度Nt：Nt = Mt / (√Vt + ϵ)

• 投影回原空间的梯度更新量G˜t：G˜t ← α·P Nt

• 更新权重Wt：Wt ← Wt-1 + η·G˜t，其中η为学习率。

• 继续增加时间步t，直到满足收敛条件。

实验结果

本文首先使用 Adam 优化器将 GaLore 与现有的低秩方法进行了比较，研究者将 GaLore 应用于 AdamW、8 bit Adam 和 Adafactor 优化器。他们采用一阶统计的 Adafactor，以避免性能下降。

同时，在大多数任务中，GaLore 都能以更少的内存占用获得比 LoRA 更高的性能。这表明，GaLore 可以作为一种全栈内存高效训练策略，用于 LLM 预训练和微调。

实操Golare！单张4090预训练7B大模型

将近4个小时才跑了171/150000,如果按照不减速，显卡不损耗的话要将近3509个小时，折合146天5个月！ 估计再过一段时间学校都要来抓我挖矿了，这是真的以时间换空间...... 然后是跑上之后显卡信息，基本上在80度左右

数据集用的是C4，下面是一些运行示例。

torchrun --standalone --nproc_per_node 1 torchrun_main.py     --model_config configs/llama_7b.json     --lr 0.005     --galore_scale 0.25     --rank 1024     --upda