全球最大开源模型再刷爆纪录!4800亿参数MoE击败Llama 3、Mixtral

新智元
主打的就是一个性价比,用比别人少好几倍的训练开销,就能得到性能差不多的模型。

就在刚刚,拥有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)的企业智能平均值与训练成本的比较

具体信息如下——

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的训练成本,被打下来了!

在以往,用LLM构建顶级企业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三阶段训练的动态数据组成

推理效率

如果希望低成本部署模型,推理效率也同样至关重要。

作为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)企业智能的平均值与活跃参数的对比

开源代码

新模型Arctic基础模型和指令微调模型代码全部开源,任何人可以将其用于研究、产品、原型当中。

项目地址: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上下文窗口。

团队介绍

Snowflake的CEO,是Sridhar Ramaswamy,是前谷歌高级副总裁。

在谷歌工作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。

他本人的研究兴趣是机器学习与系统的交叉领域。

本文作者:新智元,本文来源:新智元,原文标题:《全球最大开源模型再刷爆纪录!4800亿参数MoE击败Llama 3、Mixtral》

风险提示及免责条款
市场有风险,投资需谨慎。本文不构成个人投资建议,也未考虑到个别用户特殊的投资目标、财务状况或需要。用户应考虑本文中的任何意见、观点或结论是否符合其特定状况。据此投资,责任自负。