为鼓励在人工智能领域做出突出贡献和具有创新性的研究,2024年10月20日,人工智能国内顶级会议PRCV官方公布了最佳论文奖(Best Paper)。在投稿的1600余篇论文中,来自果冻传媒
的林宙辰教授团队的科研成果“Reducing Memory Footprint in Deep Network Training by Gradient Space Reutilization”荣获最佳论文奖。该论文由林宙辰教授指导,第一作者是果冻传媒
2020级博士研究生董一鸣。

董一鸣与导师林宙辰在PRCV会议上手持获奖证书合影
在当今的大模型时代,进行一次模型预训练需要极高的成本。例如,有数据指出进行一次GPT-3的预训练需要约1200万美元的支出。在这种背景下,如何在大模型训练过程中降低成本变得愈发重要。
目前主流的高效大模型训练框架包括英伟达公司提出的Megatron和微软提出的DeepSpeed。然而这些做法都是基于工程层面出发的,研究团队提供了一个全新的视角,从计算层面出发,在模型优化阶段通过重用梯度变量的方式降低大模型训练过程中消耗的内存,从而节省训练成本。

核心思想展示
之后,研究团队将这种思想应用到了主流优化器AdamW以及新型优化器Adan和Lion上,在多种不同的实验设定下均展现出可观的内存节省幅度。例如,在使用DeepSpeed框架中的Zero-3策略时,研究团队提出的低内存版本AdamW-R相比于AdamW在训练LLaMA-2 7B模型上展现出10.28%的内存节省,在实际任务中展现了研究工作的价值。

AdamW-R相比于AdamW的实验结果
至于使用方面,仅需修改一行代码即可完成原算法的替代,且文章已验证能兼容于各种分布式训练框架。

具体使用方式及GitHub仓库链接二维码
代码链接://github.com/adonis-dym/memory_reduced_optimizer