[论文评析]Full parameter fine-tuning for large language models with limited resources,ArXiv 2023

---

论文信息

题目: Full parameter fine-tuning for large language models with limited resources;
发表: ArXiv 2023;
作者:Kai Lv, Yuqing Yang, Tengxiao Liu, Qinghui Gao, Qipeng Guo, Xipeng Qiu;

来自复旦大学邱锡鹏老师的工作.

背景

众所周知, 训练LLM需要的计算资源非常庞大. 目前的微调主要侧重于参数高效的Fine-tune, 即对LLM增加/微调一小部分参数, 很少有相关工作聚焦于优先资源下的LLM全量微调.

简单来说, 这盘论文提出了一种新的优化器Low memory Optimizer (LOMO), 能够在极其有限的计算资源下实现对LLM的全量微调. 与标准的 与标准方法（DeepSpeed解决方案）相比，将内存使用量减少到10.8%。

方法

方法的大致原理如下图所示, 其中P表示模型的参数, P1, P2, P3表示一个模型的三部分参数, G1, G2, G3为对应的梯度信息.
通过对比,可以发现在更新模型参数:

(1) 经典的SGD主要分两步, 第一步:先把模型所有参数的梯度计算出来, 然后第二部: 再利用梯度信息更新模型的所有参数;

(2) 对于LOMO, 这两步操作的粒度更细: 依次对每一部分参数执行上述操作, 文章中称之为"Fuse the gradient computation and parameters update", 实际上就是这么回事.

这样一来, 梯度存储的复杂度就从O(n)变为了O(1).

算法伪码

通过上述的分析, 相信大家对方法有了比较清晰的认识. 这里附上伪代码.
大家可以看到:

每一部分参数实际上对应一层的参数. 在更新参数的时候实际上是按照自顶向下的方式逐层进行gradient computation & parameters update, 也就是所谓的LOMO, 哈哈.

存储使用对比分析

可以看到,在采用SGD和AdamW的方式中, Optimizer State和Gradient占据了大量的memory, 而采用LOMO后, Optimizer State占用的memory为0, 而Gradient占用的memory大大降低.

下面是数值结果对比,

达到的效果

可以看到:

(1) 对于7B的LLM, 采用LOMO后, 只用1块RTX 3090 (1/8的硬件资源)就可以达到更低的Memory峰值, 更高的Throughput;

(2) 对于13B的LLM, 采用LOMO后, 只用2块RTX 3090 (1/4的硬件资源)就可以达到同等的Memory峰值, 更高的Throughput;

(3) 对于30B的 LLM, 采用LOMO后, 只用4块RTX 3090 就可以正常训练;

(4) 对于65B的LLM, 采用LOMO后, 只用8块RTX 3090就可以正常训练.

其他

关于LOMO的理论分析以及更多的实验对比, 请阅读原文.

总结

方法超级简单, 看样子挺好用.

References
1.Full parameter fine-tuning for large language models with limited resources