10万美元训出Llama-2级大模型!全华人打造新型MoE,贾扬清SD前CEO围观

发布于 2024-4-7 09:58
浏览
0收藏

“只需”10万美元,训练Llama-2级别的大模型。


尺寸更小但性能不减的MoE模型来了:


它叫JetMoE,来自MIT、普林斯顿等研究机构。


性能妥妥超过同等规模的Llama-2。

10万美元训出Llama-2级大模型!全华人打造新型MoE,贾扬清SD前CEO围观-AI.x社区

△贾扬清转发

要知道,后者可是数十亿美元级别的投入成本。

10万美元训出Llama-2级大模型!全华人打造新型MoE,贾扬清SD前CEO围观-AI.x社区

JetMoE发布即完全开源,且学术界友好:仅使用公开数据集和开源代码,用消费级GPU就能进行微调。


不得说,大模型的打造成本,真的比人们想的要便宜更多了。


Ps. Stable Diffusion前老板Emad也点了赞:

10万美元训出Llama-2级大模型!全华人打造新型MoE,贾扬清SD前CEO围观-AI.x社区

10万美刀实现Llama-2性能

JetMoE启发于ModuleFormer的稀疏激活架构。


(ModuleFormer,一种基于稀疏专家混合(SMoE)的模块化架构,可提高大模型效率和灵活性,去年6月提出)它的注意力层中仍然使用了MoE:


80亿参数的JetMoE一共有24个区块,每块包含2个MoE层,分别是注意力头混合 (MoA) 和MLP专家混合 (MoE)。


每个MoA和MoE层又有8个专家,每次输入token激活2个。

10万美元训出Llama-2级大模型!全华人打造新型MoE,贾扬清SD前CEO围观-AI.x社区

JetMoE-8B使用公开数据集中的1.25T token进行训练,学习率5.0 x 10-4,全局batch size为4M token。


具体训练方案遵循MiniCPM(来自面壁智能,2B模型就能赶超Mistral-7B)的思路,共包含两阶段


第一阶段使用线性预热的恒定学习率,用来自大规模开源预训练数据集的1万亿个token进行训练,这些数据集包括RefinedWeb、Pile、Github data等等。

10万美元训出Llama-2级大模型!全华人打造新型MoE,贾扬清SD前CEO围观-AI.x社区

第二阶段则使用指数学习率衰减,用2500亿个token训练来自第一阶段数据集和超高质量开源数据集的token。

10万美元训出Llama-2级大模型!全华人打造新型MoE,贾扬清SD前CEO围观-AI.x社区

最终,团队使用96×H100的GPU集群,花费2周时间、约8万美元搞定JetMoE-8B。

更多技术细节将在不久后发布的技术报告上揭露。


而在推理过程中,由于JetMoE-8B仅具有22亿个激活参数,因此计算成本大大降低——

同时,它还收获了不错的性能表现。


如下图所示:

JetMoE-8B在8个评测基准上获得了5个sota(包括大模型竞技场Open LLMLeaderboard),超过LLaMA-13B、LLaMA2-7B和DeepseekMoE-16B。

10万美元训出Llama-2级大模型!全华人打造新型MoE,贾扬清SD前CEO围观-AI.x社区

在MT-Bench基准上得分6.681,也超过了130亿参数的LLaMA2、Vicuna等模型。

10万美元训出Llama-2级大模型!全华人打造新型MoE,贾扬清SD前CEO围观-AI.x社区

作者介绍

JetMoE一共4位作者,分别是:


  • Yikang Shen

MIT-IBM Watson Lab研究员,研究方向NLP。


本硕毕业于北航,博士经历于Yoshua Bengio创办的Mila研究机构。


  • 国振 (Gavin Guo)

MIT博士在读, 研究方向为3D成像的数据高效机器学习。


UC伯克利本科毕业,去年夏天作为学生研究员加入MIT-IBM Watson Lab,导师为Yikang Shen等人。


  • 蔡天乐

普林斯顿博士在读生,本科毕业于北大应用数学和计算机科学,​​目前也是Together.ai​​ 的兼职研究员,与Tri Dao合作。


  • Zengyi Qin

MIT博士在读,同时在创业,MyShell的AI研发主管。

这家公司刚刚融资了1100万美元,投资者包括Transformer的作者。

10万美元训出Llama-2级大模型!全华人打造新型MoE,贾扬清SD前CEO围观-AI.x社区

传送门:​​​​​https://github.com/myshell-ai/JetMoE​


本文转自 量子位 ,作者:量子位


原文链接:​​https://mp.weixin.qq.com/s/98TmAe_c4H64RTZXIG5yfg​

收藏
回复
举报
回复
相关推荐