低显存高效训练大模型:Llama 3-8B在单卡24GB显存上全参数微调的创新方法

原文标题:24GB单卡全量微调Llama 3-8B,仅需添加一行代码

原文作者:机器之心

冷月清谈:

香港中文大学(深圳)的研究人员开发了一种新的算法BAdam,使得大型语言模型Llama 3-8b可以在单张24GB显存的RTX 3090显卡上进行全参数微调,大幅降低了显存使用量。这一成就通过结合块坐标下降算法与大模型优化实现,显存开销仅为传统方法的六分之一。BAdam不仅在内存占用上有优势,还在损失收敛速率和迭代时间上超过其他优化算法如LoRA。此外,采用BAdam微调的Llama模型在语言任务基准测试中也表现出色,显示其实际应用潜力。该研究不仅提供了显存优化的有效方案,也对AI领域中大模型的训练提供了新思路。

怜星夜思:

1、如何评估BAdam算法在其他类型的深度学习模型上的通用性和效果?
2、显存限制常常是AI模型训练的瓶颈,除了优化算法,还有哪些其他方式可以有效解决或缓解这一问题?
3、BAdam算法能否实际帮助缩短大规模语言模型的商业应用部署周期?

原文内容

图片


ChatGPT问世以来,大型语言模型在各个领域引起了广泛兴趣,并催生了基于语言模型的应用,包括但不限于自动文本生成、信息检索、智能助理、聊天机器人以及智能教育系统等。这些应用的表现和效果往往取决于模型本身的对话能力、逻辑推理能力以及上下文理解能力等核心特征。在实际应用中,为了满足不同领域对模型能力的个性化需求,研究人员通常会基于预训练的大型语言模型进行微调,以适应特定任务的要求。

然而,微调具有一定规模的语言模型需要大量的计算资源,其中显存容量往往成为主要限制因素。主流的优化算法如Adam在训练过程中需要存储模型参数、梯度信息以及优化器状态。举例来说,当训练一个拥有70亿个参数的模型时,以上参数将占用超过120GB的显卡内存。然而,主流消费级显卡如RTX 3090/4090仅有24GB的显存,而AI训练专用显卡A100的显存也仅有80GB。因此,训练规模较大的模型不可避免地对计算资源提出了巨大需求,而且由于需要多块显卡并行训练,也带来了工程实现上的诸多挑战。

在本研究中,来自香港中文大学(深圳)的研究者通过将传统的块坐标下降算法与大模型优化结合,提出BAdam算法,将内存开销大幅降至原来的约六分之一,实现了在一张24GB显存的RTX 3090上全参数微调Llama 2-7b和Llama 3-8b模型。从优化角度来看,BAdam的损失函数收敛速率快于LoRA,且单次迭代所需时间约为LoRA的一半。从下游任务表现来看,由BAdam训练得到的Llama 2-7b, Llama 3-8b模型,在MT bench score上均领先同等参数量下的LoRA算法,并大幅领先基于SGD更新的LOMO算法。此外,在SUPERGLUE的基准测试中,由BAdam训练的RoBERTa-large模型在下游任务性能上与使用Adam进行全参数微调的模型相当,显著优于LoRA 。

为了更好的帮助大家了解这项研究,机器之心最新一期线上分享邀请到论文作者、香港中文大学(深圳)数据科学学院在读博士生罗琪竣,为大家解读他们的这项工作 。

分享主题:24GB单卡高效训练Llama 3-8B,仅需添加一行代码

嘉宾介绍:琪竣为香港中文大学(深圳)数据科学学院在读博士生,导师为李肖教授。本科毕业于香港中文大学(深圳)计算机专业,研究兴趣包含机器学习、优化理论,目前的研究重点为大语言模型的优化和泛化。

分享摘要:本次分享将概述当前参数高效微调算法及其局限性,进而引出BAdam算法设计并分析其算法特性,包含显存占用、单次迭代所需运算、以及收敛性质,随后将基于以上分析重点介绍BAdam算法实现技巧,以及超参数选取策略。
  • 论文链接:https://arxiv.org/abs/2404.02827

  • 代码链接:https://github.com/Ledzy/BAdam

直播间:关注机器之心机动组视频号,立即预约直播。

机动组技术交流群:添加小助手好友备注「研究方向」,获取行业一手资讯,欢迎大家进群聊聊。



机器之心 · 机动组
机动组聚焦于学术研究、工程实践与产业应用,筛选前沿、专业、实用内容,不定期组织学术研讨、技术交流与实战分享等欢迎所有 AI 领域技术从业者关注


点击阅读原文,直达机动组官网,查看往期回顾。

除了像BAdam这样的算法优化,还可考虑硬件升级,比如使用具有更高显存的GPU。此外,模型剪枝、量化压缩等技术也可以减少模型需要的资源,从而在现有硬件限制下提高模型性能。

噢,技术总是有更多花招的!例如,可以使用模型并行化,将大模型分割到多个GPU上进行训练,或者尝试使用更有效的数据加载和缓存策略,减少每次训练需要加载的数据量,以此来降低对显存的需求。

哎呀,别忘了,AI领域变化快,经验法则可能今天有效,明天就过时了。咱们如果真想深入探讨BAdam的威力,何不操刀一下,将它植入不同任务,比如图像识别或语音处理,看看会发生什么神奇的事情。也许能打开新世界的大门,也许能找到需要改进的方向,说走就走的测验怎么样?

你知道吗?服务器虚拟化技术也可帮大忙。通过虚拟化,多个虚拟机可共享物理GPU资源,调高使用效率。此外,一些新兴的显存交换算法允许在GPU和CPU之间动态交换数据,这也是一种不错的策略来应对显存限制。

确实,虽然BAdam在Llama模型上表现优异,但是每个模型结构不同,参数调整需求也有所不同。可以通过将BAdam应用于例如卷积神经网络(CNN)或递归神经网络(RNN)这类结构,对比其性能和资源消耗来评估其通用性。这种实验设计可以全面评估BAdam算法的适用范围和潜在限制。

尝试把BAdam应用在不同模型上肯定是个思路,但我们还得科学地来。比如设置一系列科学实验,用相同的数据集,分别用BAdam和其他优化算法在不同模型上训练,然后比较他们在性能、速度和资源使用等方面的表现。通过这样的对比测试,我们才能得出是否应广泛采纳BAdam的结论。

理论上是的。通过减少显存需求和加快迭代速度,BAdam可以使模型训练更加高效,从而加速模型从研究到实际应用的转化过程,缩短部署周期。这对于追求快速迭代和时间优先的商业环境来说非常有益。

说到缩短部署周期,咱们不能单看算法哟,得考虑整个系统的配合。比如自动化部署工具、持续集成和持续部署(CI/CD)系统等,这些都能有效缩短部署时间。当然,BAdam提高的训练效率是助推器,没有错。

这个要看具体情况,但BAdam优化的确提供了可能。通过减轻硬件限制和提高训练效率,公司能更快地迭代模型,测试新算法。当然,商业部署还涉及其他外部因素,如市场需求、产品定位等,这些都会影响部署周期。