Redian新闻
>
32卡176%训练加速,开源大模型训练框架Megatron-LLaMA来了

32卡176%训练加速,开源大模型训练框架Megatron-LLaMA来了

公众号新闻
机器之心发布
机器之心编辑部

9 月 12 日,淘天集团联合爱橙科技正式对外开源大模型训练框架 ——Megatron-LLaMA,旨在让技术开发者们能够更方便的提升大语言模型训练性能,降低训练成本,并且保持和 LLaMA 社区的兼容性。测试显示,在 32 卡训练上,相比 HuggingFace 上直接获得的代码版本,Megatron-LLaMA 能够取得 176% 的加速;在大规模的训练上,Megatron-LLaMA 相比较 32 卡拥有几乎线性的扩展性,而且对网络不稳定表现出高容忍度。目前 Megatron-LLaMA 已在开源社区上线。

开源地址:https://github.com/alibaba/Megatron-LLaMA

大语言模型的卓越表现一次又一次地超出了人们的想象。过去几个月,LLaMA 及 LLaMA2 向开源社区的全面放开,使想训练一个自己的大语言模型的人们多了一种很好的选择。在开源社区中,已有非常多的基于 LLaMA 产出的模型,包括进行续训 / SFT(如 Alpaca、Vicuna、WizardLM、Platypus、StableBegula、Orca、OpenBuddy、Linly、Ziya 等)和从零开始训练(Baichuan、QWen、InternLM、OpenLLaMA)的工作。这些工作不仅在目前各种大模型能力客观评测榜单上表现卓越,同时展现了在长文本理解、长文本生成、代码编写、数学求解等实用场景的优秀性能。更进一步,业界出现了很多有意思的产品,例如 LLaMA 结合 Whisper 的语音聊天机器人、LLaMA 结合 Stable Diffusion 的绘画软件、医学 / 法律领域的辅助咨询机器人等。

尽管从 HuggingFace 上可以拿到 LLaMA 的模型代码,但用自己的数据训一个 LLaMA 模型对个人用户或中小型组织并不是一件低成本且简单的工作。大模型的体积和数据的规模,使得在普通的计算资源上无法完成有效的训练,算力和成本成为严重的瓶颈。Megatron-LM 社区的用户在这方面的诉求非常急迫。

 
淘天集团和爱橙科技在大模型应用上有着非常广阔应用场景,在大模型的高效训练上进行了非常多的投入。LLaMA 的问世,在数据处理、模型设计、微调及强化学习反馈调整等方面都给予了包括淘天集团和爱橙科技在内的许多公司非常多的启示,也助力业务应用场景取得了新的突破。因此,为了回馈整个 LLaMA 开源社区、促进中文预训练大模型开源社区的发展,让开发者们能够更方便地提升大语言模型的训练性能,降低训练成本,淘天集团联合爱橙科技将部分内部优化技术开源,发布 Megatron-LLaMA,期望与每一位合作伙伴共建 Megatron 及 LLaMA 生态。

Megatron-LLaMA 提供了一套标准的 Megatron-LM 实现的 LLaMA,并提供了与 HuggingFace 格式自由切换的工具,方便与社区生态工具兼容。Megatron-LLaMA 重新设计了 Megatron-LM 的反向流程,使得无论在节点数较少需要开较大梯度聚合(Gradient Accumulation, GA)、或是节点数较多必须使用小 GA 的场景,都能够取得卓越的训练性能。

  • 在 32 卡训练上,相比 HuggingFace 上直接获得的代码版本,Megatron-LLaMA 能够取得 176% 的加速即便是采用 DeepSpeed 及 FlashAttention 优化过的版本,Megatron-LLaMA 仍然能减少至少 19% 的训练时间。
  • 在大规模的训练上,Megatron-LLaMA 相比较 32 卡拥有着几乎线性的扩展性例如使用 512 张 A100 复现 LLaMA-13B 的训练,Megatron-LLaMA 的反向机制相对于原生 Megatron-LM 的 DistributedOptimizer 能够节约至少两天的时间,且没有任何精度损失。
  • Megatron-LLaMA 对网络不稳定表现出高容忍度。即便是在现在性价比较高的 4x200Gbps 通信带宽的 8xA100-80GB 训练集群(这种环境通常是混部环境,网络只能使用一半的带宽,网络带宽是严重的瓶颈,但租用价格相对低廉)上,Megatron-LLaMA 仍然能取得 0.85 的线性扩展能力,然而在这个指标上 Megatron-LM 仅能达到不足 0.7。

Megatron-LM 技术带来的高性能 LLaMA 训练机会

LLaMA 是目前大语言模型开源社区中一项重要工作。LLaMA 在 LLM 的结构中引入了 BPE 字符编码、RoPE 位置编码、SwiGLU 激活函数、RMSNorm 正则化以及 Untied Embedding 等优化技术,在许多客观和主观评测中取得了卓越的效果。LLaMA 提供了 7B、13B、30B、65B/70B 的版本,适用于各类大模型需求的场景,也受到广大开发者的青睐。同诸多开源大模型一样,由于官方只提供了推理版的代码,如何以最低成本开展高效训练,并没有一个标准的范式。

Megatron-LM 是一种优雅的高性能训练解决方案。Megatron-LM 中提供了张量并行(Tensor Parallel,TP,把大乘法分配到多张卡并行计算)、流水线并行(Pipeline Parallel,PP,把模型不同层分配到不同卡处理)、序列并行(Sequence Parallel, SP,序列的不同部分由不同卡处理,节约显存)、DistributedOptimizer 优化(类似 DeepSpeed Zero Stage-2,切分梯度和优化器参数至所有计算节点)等技术,能够显著减少显存占用并提升 GPU 利用率。Megatron-LM 运营着一个活跃的开源社区,持续有新的优化技术、功能设计合并进框架中。

然而,基于 Megatron-LM 进行开发并不简单,在昂贵的多卡机上调试及功能性验证更是十分昂贵的。Megatron-LLaMA 首先提供了一套基于 Megatron-LM 框架实现的 LLaMA 训练代码,支持各种规模的模型版本,并且可以很简单地适配支持 LLaMA 的各类变种,包括对 HuggingFace 格式的 Tokenizer 的直接支持。于是,Megatron-LLaMA 可以很便捷地应用在已有的离线训练链路中,无需进行过多的适配。在中小规模训练 / 微调 LLaMA-7b 和 LLaMA-13b 的场景,Megatron-LLaMA 能够轻松达到业界领先的 54% 及以上的硬件利用率(MFU)。

Megatron-LLaMA 的反向流程优化
 

图示:DeepSpeed ZeRO Stage-2

DeepSpeed ZeRO 是微软推出的一套分布式训练框架,其中提出的技术对很多后来的框架都有非常深远的影响。DeepSpeed ZeRO Stage-2(后文简称 ZeRO-2)是该框架中一项节约显存占用且不增加额外计算量和通信量的技术。如上图所示,由于计算需要,每个 Rank 都需要拥有全部的参数。但对于优化器状态而言,每个 Rank 只负责其中的一部分即可,不必所有 Rank 同时执行完全重复的操作。于是 ZeRO-2 提出将优化器状态均匀地切分在每个 Rank 上(注意,这里并不需要保证每个变量被均分或完整保留在某个 Rank 上),每个 Rank 在训练进程中只负责对应部分的优化器状态和模型参数的更新。在这种设定下,梯度也可以按此方式进行切分。默认情况下,ZeRO-2 在反向时在所有 Rank 间使用 Reduce 方式聚合梯度,而后每个 Rank 只需要保留自身所负责的参数的部分,既消除了冗余的重复计算,又降低了显存占用。

Megatron-LM DistributedOptimizer
 


原生 Megatron-LM 通过 DistributedOptimizer 实现了类似 ZeRO-2 的梯度和优化器状态切分,以减少训练中的显存占用。如上图所示,DistributedOptimizer 在每次获得预设的梯度聚合过的所有梯度后,通过 ReduceScatter 算子,将之前累积的全部梯度分发到不同的 Rank。每个 Rank 只获得自己需要处理的部分梯度,而后进行优化器状态的更新和对应参数的更新。最后各个 Rank 通过 AllGather 的方式从其他节点上获取更新过的参数,最终取得全部的参数。实际训练的结果显示,Megatron-LM 的梯度和参数通信与其他计算串行进行,对于大规模预训练任务,为了保证总批数据大小不变,通常无法开启较大的 GA。于是通信占比会伴随机器增加上升,这时候串行通信的特点导致扩展性很弱。在社区内,这方面的需求也很迫切。
 


Megatron-LLaMA OverlappedDistributedOptimizer

为了解决这一问题,Megatron-LLaMA 改进了原生 Megatron-LM 的 DistributedOptimizer,使其梯度通信的算子能够可以和计算相并行。特别的,相比于 ZeRO 的实现,Megatron-LLaMA 在并行的前提下,通过巧妙的优化优化器分区策略,使用了更具有具有扩展性的集合通信方式来提升扩展性。OverlappedDistributedOptimizer 的主要设计保证了如下几点:a) 单一集合通信算子数据量足够大,充分利用通信带宽;b) 新切分方式所需通信数据量应等于数据并行所需的最小通信数据量;c) 完整参数或梯度与切分后的参数或梯度的转化过程中,不能引入过多显存拷贝。

 
具体而言,Megatron-LLaMA 改进了 DistributedOptimizer 的机制,提出了 OverlappedDistributedOptimizer,用于结合新的切分方式优化训练中的反向流程。如上图所示,在 OverlappedDistributedOptimizer 初始化时,会预先给所有参数分配其所属的 Bucket。Bucket 中的参数是完整的,一个参数仅属于一个 Bucket,一个 Bucket 中可能有多个参数。逻辑上,每个 Bucket 将会被连续等分成 P(P 为数据并行组的数量)等份,数据并行组中的每个 Rank 负责其中的一份。

Bucket 被放置在一个本地队列(Local grad bucket queue)中,从而保证通信顺序。在训练计算的同时,数据并行组间以 Bucket 为单位,通过集合通讯交换各自需要的梯度。Megatron-LLaMA 中 Bucket 的实现尽可能采用了地址索引,只在有需要值更改时才新分配空间,避免了显存浪费。

上述的设计,再结合大量的工程优化,使得在大规模训练时,Megatron-LLaMA 可以很充分地使用硬件,实现了比原生 Megatron-LM 更好的加速。从32张A100卡扩展到512张A100卡的训练,Megatron-LLaMA在常用混部的网络环境中仍然能够取得0.85的扩展比。

Megatron-LLaMA 的未来计划

Megatron-LLaMA 是由淘天集团和爱橙科技共同开源并提供后续维护支持的训练框架,在内部已有广泛的应用。随着越来越多的开发者涌入 LLaMA 的开源社区并贡献可以相互借鉴的经验,相信未来在训练框架层面会有更多的挑战和机会。Megatron-LLaMA 将会紧密关注社区的发展,并与广大开发者共同推进以下方向:

  • 自适应最优配置选择
  • 更多模型结构或局部设计改动的支持
  • 在更多不同类硬件环境下的极致性能训练解决方案

项目地址:https://github.com/alibaba/Megatron-LLaMA

© THE END 

转载请联系本公众号获得授权

投稿或寻求报道:[email protected]

微信扫码关注该文公众号作者

戳这里提交新闻线索和高质量文章给我们。
相关阅读
国内开源第一!百川开源大模型Baichuan-13B评测揭晓|SuperCLUE减少内存消耗、降低大模型训练成本,ACL杰出论文作者揭秘CAME优化器Erklärung zur Zusammenarbeit1800亿参数,世界顶级开源大模型Falcon官宣!碾压LLaMA 2,性能直逼GPT-4全面超越LLaMA2,月下载量超三百万,国产开源大模型如何成为新晋顶流?Pride 里的男性自卑和自尊让大模型自主探索开放世界,北大&智源提出训练框架LLaMA-RiderACL 2023 | DaMSTF: 面向领域适应的领域对抗性学习增强自我训练框架《一念如初》&《牵手爱》UC伯克利团队开源MemGPT大模型上下文内存管理方案;AgentLM、多模态Fuyu-8B、数学LLEMMA等专用大模型开源GitHub 1.1万星,模拟软件开发流程,开源框架MetaGPT爆火AI赛道陷入内卷,开源大模型纷纷商用免费持续数据开源,智源发布超3亿对面向中英文语义向量模型训练数据集开源大模型将“杀死”闭源模型?1800亿参数,性能碾压Llama 2,世界最强开源大模型Falcon 180B发布微软 Office copilot 定价每月 30 美元;腾讯否认推公务员版微信;Meta 发布新一代开源大模型 | 极客早知道​SIGIR 2023 | 用于序列推荐的多兴趣预训练框架20B跨级挑战70B性能!国产开源大模型打造大模型新标杆650亿参数大模型预训练方案开源可商用!LLaMA训练加速38%,来自明星开源项目32k上下文可商用!羊驼进化成长颈鹿,“开源大模型之最”开源大模型重击OpenAI!小扎放出LLaMA2炸裂科技圈,联手微软高通冲击市场格局Alluxio助力AI大模型训练一次通过率73%,开源代码大模型WizardCoder超越最新GPT-4以外所有闭/开源模型开源大模型FLM-101B:训练成本最低的超100B参数大模型最近一段时间我惨遭歧视!击败Llama 2,抗衡GPT-3.5,Stability AI新模型登顶开源大模型排行榜跑步者眼中的夕阳 (The Sunset in A Runner’s Eyes)LLM大模型训练Trick系列之拒绝采样nǚ hóng?nǚ gōng700亿参数LLaMA2训练加速195%,开箱即用的8到512卡LLaMA2全流程方案来了!旧文明的社会组织(第三章摘要)首个可商用的32k上下文开源大模型「长颈鹿」来了,两万词长文不在话下英伟达推出大模型加速包;Meta将训练对标GPT-4的大模型;ChatGPT等成费水大户丨AIGC大事日报大型语言模型训练加速专场在线研讨会上线,三位专家将主讲AI服务器、训练框架和云上高性能存储架构阿里云中标9亿AI算力大单;浙大校友开源LLaVA-1.5多模态大模型;Meta因开源模型遭抗议丨AIGC大事日报
logo
联系我们隐私协议©2024 redian.news
Redian新闻
Redian.news刊载任何文章,不代表同意其说法或描述,仅为提供更多信息,也不构成任何建议。文章信息的合法性及真实性由其作者负责,与Redian.news及其运营公司无关。欢迎投稿,如发现稿件侵权,或作者不愿在本网发表文章,请版权拥有者通知本网处理。