鉴于大家对GaLore比较感兴趣,我今天试着结合论文做一个更深入的解读:
GaLore的全称是"Gradient Low-Rank Projection",翻译过来就是"梯度低秩投影"。它的核心思想是通过降低优化器状态的秩,来大幅减少内存占用。
在训练大模型时,我们需要存储三类数据:模型权重、优化器状态(如动量、梯度等)和中间激活值。而优化器状态往往占据了最大的内存空间。以70亿参数的模型为例,模型权重约占14GB,中间激活值约占2GB,但优化器状态却高达42GB!
GaLore的目标就是在不影响训练效果的前提下,最大限度地压缩优化器状态的内存占用。它用到了一些很巧妙的数学技巧,比如低秩分解,将原本庞大的矩阵近似表示为两个更小的矩阵的乘积。这就像是给矩阵"瘦身",把多余的"赘肉"都去掉,只保留最精华的部分。
在实验中,GaLore对70亿参数的LLaMA模型进行了测试。结果显示,它成功将优化器状态的内存占用从42GB降到了14.5GB,降幅超过65%!再加上一些其他的内存优化技术,最终将总内存消耗控制在了21.3GB,比原来的58GB低了很多。
这意味着我们现在可以在家用的24GB显卡上,轻松训练70亿参数的大模型了!这在此前是根本不可想象的。
与其他一些内存优化方法相比,GaLore的优势在于它既支持从头训练,又能显著降低内存占用。像Colossal-AI是通过将模型切分到多个设备上来分担压力,但总内存需求其实没变;LoRA是在预训练模型上叠加一个小矩阵来微调,内存占用是小了,但只适合微调,不能从头训练。
GaLore的出现,为大模型训练开辟了一条全新的道路。它让每个人都有机会参与到大模型的开发中来,用自己的创意和智慧去塑造心目中最理想的AI助手。
