全球最大开源模型再刷爆纪录!4800亿参数MoE击败Llama 3、Mixtral
新智元报道
新智元报道
【新智元导读】最大开源模型,再次刷爆纪录!Snowflake的Arctic,以128位专家和4800亿参数,成为迄今最大的开源模型。它的特点,是又大又稀疏,因此计算资源只用了不到Llama 3 8B的一半,就达到了相同的性能指标。
就在刚刚,拥有128位专家和4800亿参数的Arctic,成功登上了迄今最大开源MoE模型的宝座。
它基于全新的Dense-MoE架构设计,由一个10B的稠密Tranformer模型和128×3.66B的MoE MLP组成,并在3.5万亿个token上进行了训练。
不仅如此,作为一个比「开源」更「开源」的模型,团队甚至把训练数据的处理方法也全给公开了。
Arctic的的两个特点,一个是大,另一个就是非常稀疏。
好处就在于,这种架构让你可以用比别人少好几倍的训练开销,就能得到性能差不多的模型。
也就是说,与其他使用类似计算预算训练的开源模型相比,Arctic的性能更加优异。
比起Llama 3 8B和Llama 2 70B,Arctic所用的训练计算资源不到它们的一半,评估指标却取得了相当的分数!
图1 编码(HumanEval+和MBPP+)、SQL生成(Spider) 和指令遵循(IFEval)的企业智能平均值与训练成本的比较
具体信息如下——
480B参数,生成期间17B处于活跃状态;
128位专家,有2位在生成期间活跃;
Instruct & Base版本发布;
专注于企业任务(代码、SQL、推理、跟踪);
在Apache 2.0下发布;
FP16精度下约为900GB内存,INT4精度下约为240GB内存
使用DeepSpeed-MoE训练。
主打的就是一个性价比
企业智能指标,是对企业客户至关重要的技能,包括包括编码(HumanEval+和MBPP+)、SQL生成(Spider)和指令遵循(IFEval)。
同时,团队也采用了业界常用的评估LLM的学术基准,包括世界知识、常识推理和数学能力。
可以看到,Arctic在多项企业智能指标中,都超越了Mixtral 8×7B等开源对手。
在计算类别中,它实现了顶级性能,甚至和使用更高计算预算训练的模型,都有的一拼。
在学术基准上,它的表现也不差。
在测评中,团队发现了一件有意思的事。
MMLU等世界知识指标,是人们常用的学术基准测试。而随着高质量网络和STEM数据的增加,MMLU的得分会随着训练FLOPS的增加而提高。
但是,Arctic的目标之一,是在保持较小训练预算的同时优化训练效率,因此,跟其他模型相比,Arctic在MMLU上的得分较低,也实属情理之中。
由此,如果训练计算预算高于Arctic的训练,MMLU性能就将超越Arctic。
当然,MMLU世界知识的性能,并不一定和团队所关注的企业智能直接相关。
表3 Arctic与DBRX、Llama 3 8B、Llama 3 70B、Mixtral 8x7B、Mixtral 8x22B的对比
企业级AI的训练成本,被打下来了!
通常,花费的成本高达数千万甚至数亿美元,这一成本是惊人的。
如何解决有效训练和推理的限制?Snowflake AI团队的研究者一直在做这方面的努力,团队成员过去曾开源了ZeRO、DeepSpeed、PagedAttention/vLLM和LLM360等系统,显著降低了LLM训练和推理的成本。
而今天推出的Arctic,在SQL生成、编码和遵循基准指令等企业任务上,表现非常出色。
它为具有成本效益的训练设定了新的基准,用户可以以极低的成本,就能创建满足企业需求的高质量定制模型。
Arctic也是一个真正的开放模型,在Apache 2.0许可下,提供对权重和代码的无限制访问。
从今天开始,Snowflake Arctic就可以从Hugging Face上获取了。
计算资源仅用一半,表现却和Llama 3 8B相当
团队发现,企业客户对AI有着一致的需求和使用场景——构建对话式SQL数据助手、代码助手和RAG聊天机器人。
为了便于评估,团队通过对编码(HumanEval+和MBPP+)、SQL生成(Spider)和指令跟随(IFEval)取平均值,将这些能力整合到「企业智能」这个单一指标中。
在开源LLM中,Arctic仅用不到200万美元(相当于不到3000个GPU周)的训练计算预算,就实现了顶级的企业智能。
更重要的是,即使与那些使用显著更高计算预算训练的模型相比,它在企业智能任务上也表现出色。
结果显示,Arctic在企业级评估指标上的表现,与Llama 3 8B和Llama 2 70B相当,甚至更优,而它所使用的训练计算资源却不到后两者的一半。
具体来说,Arctic使用的计算预算只有Llama3 70B的1/17,但在编程(HumanEval+和MBPP+)、SQL(Spider)和指令跟随(IFEval)等企业级任务上,都与其不相上下。
表1 Arctic、Llama-2 70B、DBRX和Mixtral 8x22B的模型架构和训练计算量(与活跃参数和训练token的乘积成正比)
此外,Arctic的高训练效率还意味着,Snowflake客户和整个AI社区可以以更加经济实惠的方式训练定制模型。
训练效率
为了实现如此高的训练效率,Arctic采用了独特的Dense-MoE Hybrid transformer架构。
该架构将一个10B规模的稠密Transformer模型与一个128×3.66B规模的残差MoE MLP相结合,虽然总参数量达到480B,但通过top-2 gating的方式只选择了其中17B个参数保持活跃。
Arctic的设计和训练基于以下三个关键创新:
1. 更多但精炼的专家,以及更多的专家选择
首先,DeepSpeed团队在2021年末便证明了,MoE(Mixture of Experts)可以在不增加计算成本的情况下,显著提高LLM模型的质量。
其次,模型质量的提升主要取决于MoE模型中专家的数量、总参数量以及这些专家可以组合在一起的方式和数量。
基于此,Arctic被设计为拥有480B个参数,分布在128个细粒度专家中,并使用top-2 gating选择17B个活跃参数。相比之下,最近的MoE模型使用的专家数量就要少得多了(如表2所示)。
从直观上看,Arctic利用更大的总参数量和众多专家来扩大模型容量,同时更明智地在众多精炼的专家中进行选择,并使用适度数量的活跃参数来实现资源高效的训练和推理,最终获得顶级的智能。
图2 标准MoE架构 vs. Arctic
2. 架构和系统协同设计
即便是用最强大的AI硬件,想要基于普通的MoE架构训练大量专家效率依然很低。
其原因在于,专家之间存在的全通信开销非常高昂。不过,如果能将通信与计算重叠,那么就可以极大地降低这种开销。
因此,团队在Arctic架构中将一个密集的Transformer与一个残差MoE组件(图2)相结合,从而使系统能够通过通信计算重叠来消除大部分通信开销,最终实现了极佳的训练效率。
3. 面向企业的数据课程
要在代码生成和SQL等企业指标上表现出色,需要与训练通用指标的模型截然不同的数据课程。
团队在进行了数百次小规模的对比实验后发现,常识推理等通用技能可以在开始时学习,而编码、数学和SQL等更复杂的指标可以在训练的后期有效学习。
因此,Arctic采用了三阶段课程进行训练,每个阶段的数据组成不同——
第一阶段(1T Tokens)侧重于通用技能,后两个阶段(1.5T和1T Tokens)侧重于企业级技能。
表2 Arctic三阶段训练的动态数据组成
推理效率
训练效率,只是Arctic高效的其中一个方面。
如果希望低成本部署模型,推理效率也同样至关重要。
作为MoE模型规模的飞跃,Arctic使用了比其他开源自回归模型更多的专家和参数。
因此,为了有效地在Arctic上运行推理,团队做了一些系统性的创新——
a) 在较小batch的交互式推理中(比如批大小为1),MoE模型的推理延迟受到了读取所有活跃参数所需时间的瓶颈,其中,推理是受内存带宽限制的。
在这样的批大小下,Arctic(17B活跃参数)的内存读取次数比Code-Llama 70B少4倍,比 Mixtral 8x22B(44B活动参数)少2.5倍,从而实现更快的推理性能。
为此,团队跟英伟达的TensorRT-LLM和vLLM团队展开合作,为交互式推理提供了Arctic的初步实现。
通过FP8量化,团队可以将Arctic放入单个GPU节点中。
虽然仍远未完全优化,但在批大小为1时,Arctic的吞吐量超过70+token/秒,这样就实现了有效的交互式服务。
b) 当批大小的规模显著增加,例如每次前向传递要处理数千个token时,Arctic就会从内存带宽受限转变为计算受限,此时推理的瓶颈就在于每个token的活跃参数。
在这一点上,与CodeLlama 70B和Llama 3 70B相比,Arctic的计算需求减少了4倍。
为了实现计算受限的推理和与Arctic中活跃参数数量较少相对应的高吞吐量(如下图所示),需要较大的batch size。
要实现这一点,需要有足够的KV缓存内存来支持较大的batch size,同时也需要足够的内存来存储近500B的模型参数。
面对这重重挑战,最终团队还是找到了办法。
通过使用FP8权重、分割融合和连续批处理、节点内的张量并行性以及节点间的管线并行性等系统优化组合,团队在双节点推理中,实现了这一目标。
图3 推理期间编码(HumanEval+和MBPP+)、SQL生成(Spider)和指令跟踪 (IFEval)企业智能的平均值与活跃参数的对比
开源代码
项目地址:https://github.com/Snowflake-Labs/snowflake-arctic
研究人员基于LoRA的微调的pipeline和配方(recipe),并允许在单个节点上进行高效的模型微调。
现在,Snowflake正在与英伟达TensorRT-LLM和vLLM开展合作,为Arctic模型开发初始的推理实现,并且针对批大小为1的交互式使用进行了优化。
未来,他们还将与社区合作,解决真正大型MoE更大的批大小的推理复杂性。
Cookbook:https://medium.com/snowflake/snowflake-arctic-cookbook-series-exploring-mixture-of-experts-moe-c7d6b8f14d16
另外,Arctic现使用的是4k上下文窗口进行训练,研究人员还将开发一种基于注意力下沉(attention-sinks)的滑动窗口的方法,以支持未来几周无限序列生成能力。
下一步,将会扩展到32K上下文窗口。
团队介绍
在谷歌工作15年后,他成为Neeva的联合创始人,后来Neeva被Snowflake收购。
他在印度理工学院马德拉斯分校获得计算机学士学位,并在布朗大学获得计算机博士学位。
AI团队的一把手Vivek Raghunathan,也是前谷歌副总裁。
他曾担任微软研究员,后在谷歌从事机器学习、广告基础架构等方面工作,18年开始在谷歌担任副总裁,领导YouTube团队。
随后,他和Sridhar Ramaswamy共同创办了Neeva。
Raghunathan同样也是印度理工学院的校友,不过是在孟买分校获得的学士学位。之后,他在UIUC取得了硕士和博士学位。
为了发展AI,两人把DeepSpeed团队最顶尖的几个元老都挖了过来,包括Zhewei Yao和Yuxiong He。
Zhewei Yao在UC伯克利获得博士学位,研究兴趣在于计算统计、优化和机器学习。(在此之前,他2016年曾获得上交大数学学士学位。)
他从2021年开始便加入了微软,在微软担任首席研究员和研发经理,致力于高效的大规模训练和推理。
目前,他是Snowflake的高级科学家和SDE II,同时也是Snowflake大规模预训练创始成员。
Yuxiong He在微软任职13年,是DeepSpeed的创始人之一,最近加入了Snowflake。
她曾在新加坡南阳理工大学获得了计算机工程学士学位。
团队的另一位华人大牛Aurick Qiao,去年11月刚加入Snowflake。
CMU读博期间,他曾获得Osdi 2022的最佳论文优胜奖。此前曾在微软、Dropbox工作。
曾担任Petuum CEO,以及LMNet的联合创始人。
Hao Zhang是UCSD的Halıcıoğ数据科学研究所和计算机科学与工程系的助理教授。
他曾获得了CMU计算机博士学位,师从Eric Xing。在攻读博士学位期间,他休学一段时间并在ML平台初创公司Petuum工作。
Hao Zhang在2023年联合创立了LMnet.ai,这家公司于同年11月加入了Snowflake。
他此前还共同创办了非营利组织LMSYS Org,该组织训练了此前十分流行的模型Vicuna以及发起和维护目前最重要的大语言模型评测机制:Chatbot Arena。
他本人的研究兴趣是机器学习与系统的交叉领域。
微信扫码关注该文公众号作者