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

发布网友发布时间：2024-10-22 21:28

共1个回答

热心网友时间：2024-11-02 10:17

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

而通过梯度低秩投影，GaLore可以在不牺牲模型性能的前提下显著降低内存消耗，使得在消费级GPU上训练大型语言模型成为可能。该策略的核心思想是利用权重矩阵梯度在训练过程中表现出的缓慢变化的低秩结构，而不是直接将权重矩阵近似为低秩矩阵。具体做法是计算两个投影矩阵P∈Rm×r和Q∈Rn×r，用以将梯度矩阵G投影到低秩形式P⊤GQ，从而大幅减少依赖于梯度分量统计信息的优化器状态所占用的内存。

对于链式可逆神经网络，其权重矩阵Wl在第l层的梯度Gl可以用特定形式表示，其中包含前向传播中的雅可比矩阵及其乘积，以及与输出相关的项。定理3.2给出了可逆模型梯度形式，适用于链式可逆神经网络。在每个预设的频率T时刻重新计算投影矩阵P和Q，通过在不同的低秩子空间中进行更新来逐步*近最优解。在附带的Adam优化器结合GaLore的算法描述中，每经过一定频率T的时间步长，根据当前梯度Gt重新初始化投影矩阵，并基于此进行低秩投影后的梯度更新，以实现高效、全参数学习下的内存节省。

在实验结果部分，研究者将GaLore与现有的低秩方法进行了比较，并将其应用于AdamW、8 bit Adam 和 Adafactor 优化器。在大多数任务中，GaLore都以更少的内存占用获得了比LoRA更高的性能，这表明GaLore可以作为一种全栈内存高效训练策略，用于LLM预训练和微调。

实操方面，在单张4090上预训练7B大模型需要将近4个小时，按照不减速、显卡不损耗的假设，大约需要3509个小时，即大约146天5个月的时间。显卡运行时温度保持在80度左右。数据集使用C4，下面是运行示例，每秒钟可以处理8个样例，每秒钟可以处理2000个单词，5,000步中150次更新。通过这些实验，GaLore为个人训练大模型提供了希望。