自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 领域技术从业者关注。
点击阅读原文，直达机动组官网，查看往期回顾。
继续阅读
阅读原文