资讯 人工智能
此为临时链接,仅用于文章预览,将在时失效

MoE 高效训练的 A/B 面:与魔鬼做交易,用「显存」换「性能」

作者:晓楠
2024/05/27 15:15


MoE 会成为未来大模型训练的新方向吗?

这是人们发现 MoE 架构可以用于大模型训练、推理后,发出的一声疑问。

MoE(Mixture of Experts),又称「混合专家」,本质是一种模块化的稀疏激活。怎么理解?

当前的大模型主要分为稠密(dense)模型与稀疏(sparse)模型,两者的区别主要在于模型进行计算时,被调用的参数数量,参数全部生效使用的是稠密模型,比如 OpenAI 从第一代到第三代即 GPT-1、 GPT-2、 GPT-3,以及 Meta 的 Llama 系列都是稠密模型;只使用其中一部分参数的是稀疏模型,比如基于 MoE 架构的模型,而这些被使用的参数称为「激活参数」。

具体从网络结构来看,目前主流的大模型大都是基于 Transformer 架构,由多个 Transformer Block 叠加组成,在每一个 Transformer Block 内部都会包括两层结构,一层是多头自注意力(Multi-Head Self-Attention),另一层是位置前馈神经网络(Position-wise Feed-Forward Network,FFN)。MoE 做法是对 FFN 层进行横向扩展,将其中的参数细化、组织成一个又一个的组,每个组是一个专家(expert)。

基于此,对于特定输入,MoE 只需激活少数 expert,让其参数处于激活状态,从而实现稀疏激活。而具体激活哪些 expert、哪部分参数,MoE 还会引入一个 Router(门控网络或路由),根据特定输入进行动态选择。

去年年底,Mistral AI 开源了基于 MoE 架构的模型 Mixtral 8x7B,性能表现可达到与 GPT-3.5、 Llama2 70B 相媲美的程度,而推理时只消耗了一个 13B 级别的稠密模型计算量——这一度让业内人士对 MoE 架构在大模型研究中的潜力充满了想象。

MoE 高效训练的 A/B 面:与魔鬼做交易,用「显存」换「性能」

居然有这样一个架构,可以让 AI 模型在增大参数规模的同时,计算成本远低于相同参数规模的稠密模型,实现高效训练。这在 Scaling Law 仍被奉为 LLM 第一性原理的当下,简直是「大力出奇迹」的不二工具。

紧接着,国内外大模型从业者相继推出基于 MoE 架构的大模型:

今年 1 月,MiniMax 发布基于 MoE 架构的大语言模型 abab6;

2月,昆仑万维正式发布新版 MoE 大语言模型「天工 2.0」;

3月,马斯克创办的 xAI 宣布开源基于 MoE 架构的大模型 Grok-1,新晋大模型独角兽阶跃星辰发布 Step-2 万亿参数 MoE 语言大模型预览版,阿里通义千问团队开源首个 MoE 模型——Qwen1.5-MoE-A2.7B;

4月,元象科技发布 XVERSE-MoE-A4.2B 大模型,面壁智能推出 MiniCPM-MoE-8x2B MoE 模型;

5月,DeepSeek AI 开源 MoE 语言模型 DeepSeek-V2,在一些任务表现中成为目前最强的开源 MoE 语言模型……

越来越多的玩家开始选择 MoE。为此,AI 科技评论走访了多位国内 MoE 专业人员。我们发现,MoE 的未来或尚未可知,但在当下,MoE 是大多数企业在忠于 Scaling Law 这一定律下必然要做出的选择。

1.MoE 是算力匮乏的「产物」

MoE 并非新生事物。作为一种统计学架构,MoE 早在 1997 年就被提出,后来,随着机器学习等人工智能技术的兴起,MoE 的思想也随着深度学习的发展一路演进,在2017年时,就已经有专家学者在 LSTM 模型上运用 MoE 方法并取得了成功。

而意识到 MoE 可以用于大模型的训练,很多人是从去年年底 Mistral AI 开源了首个 MoE 架构模型开始的,但对于很多研究 MoE 的学者来说,Mistral AI 的开源更是 MoE 的一次成功实践,让大家对 MoE 的感知从原先的理论到实际效果的转变。

实际早在去年 6 月份,在一次访谈中,就有人声称 OpenAI 的 GPT-4 是一个 8x220B 的 MoE 模型。只不过彼时大模型在国内才刚刚起步,大家对基础模型的训练技术不太成熟,对 MoE 感知不够,一阵喧哗之后讨论声渐渐消散,不过也正是这一次让更多业内人士开始意识到 MoE 的价值。

阶跃星辰算法专家告诉雷峰网,当初社区传言 OpenAI 的 GPT-4 是 MoE 架构时,他们并不感到意外,甚至能根据经验猜测它的大部分架构细节。

元象XVERSE技术合伙人、大模型算法负责人轩文烽也表示,在更早之前,2021 年初 Google 在 arXiv上 发表了一篇论文《Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity》后,他们就开始了解到 MoE,同年国内的阿里正在研发 M6 模型,在这方面也有跟进。「而 Mistral AI 的开源,直接让它的实际效果呈现得更明显,之后,去年 12 月元象XVERSE开始策划训练 MoE 大模型。」

诚然,MoE 作为一种模块化的稀疏激活,可以在模型总参数量不变的情况下,大幅降低计算总量,高效计算、降低成本。这是 MoE 最为显著的优势,也是为什么它作为一个并不「新鲜」的技术又在此时此刻重新焕发活力的原因所在。

现在的大模型越来越大了。Scaling Law 带来的启示是模型越大,性能越强,可是模型迭代离不开高效的算力支持,而算力的不足和匮乏,行业周知。

稀疏化包括 MoE 技术为这一问题提供了解法——它可以增大模型的参数规模,但是它的计算成本又要远低于相同参数规模的稠密模型。

面壁智能研究员宋晨阳告诉雷峰网,在目前算力,特别是国内算力普遍比较有限的情况下,MoE 架构是一个必然的趋势和选择。

面壁智能与 MoE 渊源颇深。早在 2021 年 6 月,清华大学刘知远教授牵头的「悟道·文源」就发布了千亿 MoE 大模型 CPM-2 ,是国内最早关注到这一问题的 AI 团队之一,而其中的参与成员就包括面壁智能的初始团队。因着这一自带的基因,面壁智能一直在关注和研究 MoE 和稀疏化。

MoE 高效训练的 A/B 面:与魔鬼做交易,用「显存」换「性能」

同样坚信 Scaling Law 的阶跃星辰在完成千亿模型的训练后,开始训练万亿模型。而想把模型参数扩大到万亿,MoE 是在性能、参数量、训练成本、推理成本这些维度权衡下的最佳选择。

具体来看,MoE 模型的优势主要体现为专家化、动态化、稀疏化。

一是具有更强的表现力:MoE 架构融合多个各有所长的专家模型,既有分工,又有合作,能够动态地适应不同输入样本的特点,提高模型的表达能力。

二是训练和推理速度更快:相比于同参数的稠密模型架构,MoE 是稀疏的,训练的时候计算量更小,推理的时候同时激活的参数量更小,因此能够更高效地利用计算资源。

2.路径抉择不一

目前,市场上选择 MoE 架构的大模型厂商越来越多,大家做得也都大同小异,但在架构细节上还有很多值得琢磨的地方,业界也没有达成共识,比如总共需要多少个 expert(2个到128个)、选择激活多少个 expert、需不需要有固定共享 expert、选择 expert 的时候该怎么做均衡等,各家都会不太一样。

一般来说,现在的 MoE 路径选择主要分为两种——基于一个已有模型开始训练,还是从头开始训练。

基于已有模型开始训练,是通过一个众所周知的取巧方案——upcycle(向上复用)进行的。具体来说,是将一个已经训练好的大模型的前馈神经网络权重复制若干份,每一份作为一个 expert,然后组合在一起,这样就可以把一个非 MoE 模型变成一个 MoE 模型,参数量很快就能变得特别大,但算力需求会小几个量级。

宋晨阳告诉雷峰网(公众号:雷峰网),「这种方法的最大优势在于可以节省很多成本,因为初始模型就有很多知识。」但同时,这样也存在很大的弊端,比如基于拷贝复制得到的 MoE 模型,expert 之间会有很多知识重叠,权重、参数的冗余导致对各 expert 的支持是重叠的,如果做得不好,就不能充分发挥 MoE 架构的潜力,是一个很具考验性的路径。

那如果从头开始训练的话,虽然训练起来会更难,但上述问题会有所缓解,且它的自由度更高,避免传统 MoE 架构模型的潜在问题,争取达到 MoE 架构模型的上限。

目前,大多数厂商选择的是基于已有模型进行训练,像是 Mistral AI 的 Mixtral 系列模型、千问的 MoE 模型 Qwen1.5-MoE-A2.7B、面壁智能的 MiniCPM-MoE-8x2B 等。而从论文来看,DeepSeek 选择的是从头开始训练,前段时间 DeepSeek-V2 一经开源发布,引起了非常大的反响。

阶跃星辰的做法是,一方面选择完全从头开始训练,保持万亿参数在整个训练过程中都不变,充分发挥 MoE 架构的潜力,另一方面在前期做好准备工作,从众多方案中选择更为靠谱的架构配置。

除此之外,还有一个大家会有分歧的问题,就是 MoE 架构的 expert 的力度应该是粗还是细?以及应不应该有固定共享 expert?

举例来说,Mistrial AI 的 8x7B 是一个经典的 MoE 架构,有 8 个 expert ,每个 expert 都是完全一样的形状,然后用 router 网络去从 8 个 expert 里面选择 2 个 expert。但是千问和 DeepSeek 采用了不同的做法,特别是 DeepSeek 把 expert 的粒度切得更细,同时还引入共享 expert。

MoE 高效训练的 A/B 面:与魔鬼做交易,用「显存」换「性能」

换句话说,好比原本是从 8 个 expert 中只能选 2 个,那组合数为 C8 2,但是如果把这 8 个 expert 每个再分别切成 8 个,就得到 64 个 expert,然后从 64 个 expert 中去选择 8 个,组合数就达到了 C64 8。这样一来,模型可选择的组合大大增多,模型的表达能力也就会有所增强。

宋晨阳认为,这是 DeepSeek MoE 模型的一大创新之处。

对此,轩文烽也告诉雷峰网,在 expert 的选择上,元象XVERSE 也没有使用传统的 MoE 架构做法——类 Mistral AI 的选择,每个 expert 等同于标准的 FFN——而是用了一些更细颗粒度的方法,将 expert 的设计设置为标准 FFN 的 1/ 4,而之所以做出这个选择,是在训练过程中,根据多次实验预测结果得出的结论,另外他们也在密切关注目前业界的 FFN 的配比方案。

另外,DeepSeek 还引入了共享 expert,前面提到,如果像 Mistral AI 那样从一个已有的稠密模型基于 upcycle,得到的 MoE 模型存在一个问题,即每个 expert 在初始化的时候完全一样,彼此之间不可避免会有很多知识的重叠,最后参数将是冗余的。而 DeepSeek 的策略是把每个 expert 重叠的知识用共享 expert 单独拿出来表示,共享 expert 对于每一个给定的输入 token 不经过 router 网络直接固定激活,然后再对其他剩下不共享的 expert 中,用 router 网络去选择要激活的部分。

基于此,宋晨阳认为,目前市面上的 MoE 模型又分为两个流派,一种就是类似 Mistrial AI的传统 MoE 流派,还另外就是千问和 DeepSeek MoE 这种 expert 切分更细、设置共享 expert 的架构。

至于哪种会成为主流,其实并没有具体的答案,因为一个模型最后表现有多好,架构还真不一定是最重要的事情,影响因素很多,比如数据,包括哪些类型的数据、数据如何清洗、如何配比组织,以及最终训练时,原始数据如何切分成一个又一个的 token,等等。不同的选择都将通往不同的答案。

当前,面壁智能的 MiniCPM-MoE-8x2B 采用的是类似 Mistral AI 的模型架构。考虑到未来面壁智能的最终目标是要将端侧模型做到极致,送到每一部手机、每一个用户手上。所以,当下他们也在讨论,包括做一些实验,来决定设计怎样更适合加速的 MoE 架构,从而做出最终路径选择。

3.MoE「终究是一门妥协的艺术」

当前,随着入局玩家的增多,MoE 似乎正逐渐成为一种行业共识,那么,MoE 会是未来 AI 的方向吗?

对此,众说纷纭,信奉它的人认为,它可以实现同等参数规模下算力的大幅降低,在 Scaling Law 仍未失效的当下,无疑是最优解。

但不看好的也大有人在。

今年 4 月底,Meta 官宣了 Llama 3,总共有三个版本——8B、70B 和 400B+。其中,8B 和 70B 版本已经开源,而 400B+ 版本仍在训练中。CEO 扎克伯格在访谈提到,Llama 3 400B+ 将是一个稠密模型。面对如此大的参数量,仍然选择稠密模型,或许表明,Meta 不相信稀疏化,「拒绝」了 MoE 架构。

其实,即便是那些选择了 MoE 架构的玩家,也不得不承认,虽然有着高效训练、灵活性和可解释性的优点,MoE 的局限性仍然明显,是「一门妥协的艺术」。

宋晨阳表示,MoE 的好处在于算力有限的情况下,可以实现同等参数规模下,计算量的大幅降低。但如果有足够的算力,训练一个 MoE 模型和同等参数量的稠密模型,就性能来说,一定是稠密模型更好。

也就是说,同等参数规模下,稠密模型的表现要优胜于 MoE 模型。这是因为,本质上 FFN 的表达能力是不受限的,但 MoE 模型为了稀疏激活,强行规定只能激活被切细后的一部分 FFN,相当于把模型的整个可表示空间给缩小了。

另外,无论是 MoE 还是整个稀疏化,都有一个「致命」问题,就是目前并没有一个特别有效的方法能够降低它们的存储消耗。比如对于一个 7B 模型来说,可以把它做成稀疏架构,使其计算量小于 7B,但它的内存消耗还是一个 7B 的规模,这会导致在实际将其部署在端侧的时候,出现问题。

轩文烽也表示,MoE 现在的高效训练等优势其实就是通过显存换效果。虽然 MoE 不像稠密模型那样,参数增大后训练、推理成本都会显著增加,但对于一个用 4.2B 的激活参数量达到 13B 效果的 MoE 模型来说,推理成本是低了些,但是显存的占用并没有变化。

因此,「如何优化显存方面的问题,同时保持它这种获得更高效果的性价比,可能是需要行业去研究的。」

另外,一位从事 MoE 模型训练的研究员也感慨,在实际训练中,MoE 架构模型确实存在一些挑战。

比如,训练和推理难度大,万亿的 MoE 架构需要将各专家放在不同的计算节点上,会引入额外的并行维度和网络通信,专家之间的不平衡性会拖累整个系统的训练、推理效率,极端情况下可能会浪费 50% 以上的集群算力;参数利用率不高,每个专家都必须成为一定程度上的「通才」,消耗很多的参数量学习共有的基础知识,在此基础上才能有所「专精」,因此相同参数量的 MoE 模型往往比不过传统稠密模型。

但这些从业者也纷纷向雷峰网坦言,即便如此,就当下的实际情况而言,MoE 肯定是一个非常有希望的做法、必然的趋势。

或许,在算力足够的理想状态下,没有必要非得用 MoE 架构,但在算力依旧短缺的现在,无法设想那个不存在的未来。当然,如果我们把目光放得长远些,十年或二十年后,技术持续发展,硬件和计算框架不再成为短板,到底想要使用什么样的架构训练大模型,会更自由些。

而在当下这个明确而具体的时间点,从业者能做的就是基于 MoE 架构现有的一些技术挑战,以及可优化空间,尽量进行可能性的改进,使它的训练更高效、性能更好。

比如,如何为 MoE 设计一个更好的加速框架?MoE 模型并不是一个特别好加速的模型,因为有很多 expert ,这些 expert 如何部署,是把每个 expert 部署在不同的卡上,还是把它们切片以后进行平行、分布式部署在多张卡上,等等。

再比如,目前关于 MoE 大模型似乎开始出现一个趋势,就是当大家都在争着把大模型的参数量做得越来越大时,性能却并没有因此变好,甚至没有参数量更小的模型好。一组对比是,参数量为 314B 亿的 Grok-1 与 Mistral AI 的 8x7B 模型性能相当,这是为什么?

对此,宋晨阳认为其中一个核心原因是 Mistral AI 的数据质量更高。Scaling Law 能够成立的前提,是要有充足的高质量数据作为保证。对于同样的模型架构、同样的参数量来说,数据质量比较高的那个模型,可以用比较少的数据量达到相同的效果。

另外,除了数据,另一个决定模型性能因素是训练技巧,包括批次大小等调度,都非常有讲究,但目前这些在行业中,很多已经做得非常好的头部厂商「羞于」展示,需要更多的从业者不断实践,获得更多的经验来探索更好、更优的方式。当然,也有一些厂商敢于「逆行」,如面壁智能的 MiniCPM-2B 在发布的同时,于技术报告中详细描述了如批次大小调度、学习率调度、词表大小选取等细节问题,为学界和业界提供了可参考的经验。

比如轩文烽在总结 XVERSE-MoE-A4.2B 大模型的研发经验时觉得,如果再重新来一遍,可能会做得更快些,实验计划会设计得更紧凑些。

虽然,大家都想找到一个最优的解决方案,将大模型训练、推理得又快又好,但不可否认,技术探索没有边界,没有最优、只有更优。

MoE 包括稀疏化未必代表 AI 的未来,但它提供了丰富的想象空间。比如,是不是可以朝着更为稀疏的方向走去,将 expert 切得更细,细到极限,如实现神经元级别的稀疏激活,那又将会带来怎样的效果?

这是一个非常有前景、值得探索的方向。


长按图片保存图片,分享给好友或朋友圈

MoE 高效训练的 A/B 面:与魔鬼做交易,用「显存」换「性能」

扫码查看文章

正在生成分享图...

取消
相关文章