10萬美元訓(xùn)出Llama-2級大模型!全華人打造新型MoE,賈揚清SD前CEO圍觀
“只需”10萬美元,訓(xùn)練Llama-2級別的大模型。
尺寸更小但性能不減的MoE模型來了:
它叫JetMoE,來自MIT、普林斯頓等研究機構(gòu)。
性能妥妥超過同等規(guī)模的Llama-2。
△賈揚清轉(zhuǎn)發(fā)
要知道,后者可是數(shù)十億美元級別的投入成本。
JetMoE發(fā)布即完全開源,且學(xué)術(shù)界友好:僅使用公開數(shù)據(jù)集和開源代碼,用消費級GPU就能進行微調(diào)。
不得說,大模型的打造成本,真的比人們想的要便宜更多了。
Ps. Stable Diffusion前老板Emad也點了贊:
10萬美刀實現(xiàn)Llama-2性能
JetMoE啟發(fā)于ModuleFormer的稀疏激活架構(gòu)。
(ModuleFormer,一種基于稀疏專家混合(SMoE)的模塊化架構(gòu),可提高大模型效率和靈活性,去年6月提出)
它的注意力層中仍然使用了MoE:
80億參數(shù)的JetMoE一共有24個區(qū)塊,每塊包含2個MoE層,分別是注意力頭混合 (MoA) 和MLP專家混合 (MoE)。
每個MoA和MoE層又有8個專家,每次輸入token激活2個。
JetMoE-8B使用公開數(shù)據(jù)集中的1.25T token進行訓(xùn)練,學(xué)習(xí)率5.0 x 10-4,全局batch size為4M token。
具體訓(xùn)練方案遵循MiniCPM(來自面壁智能,2B模型就能趕超Mistral-7B)的思路,共包含兩階段:
第一階段使用線性預(yù)熱的恒定學(xué)習(xí)率,用來自大規(guī)模開源預(yù)訓(xùn)練數(shù)據(jù)集的1萬億個token進行訓(xùn)練,這些數(shù)據(jù)集包括RefinedWeb、Pile、Github data等等。
第二階段則使用指數(shù)學(xué)習(xí)率衰減,用2500億個token訓(xùn)練來自第一階段數(shù)據(jù)集和超高質(zhì)量開源數(shù)據(jù)集的token。
最終,團隊使用96×H100的GPU集群,花費2周時間、約8萬美元搞定JetMoE-8B。
更多技術(shù)細節(jié)將在不久后發(fā)布的技術(shù)報告上揭露。
而在推理過程中,由于JetMoE-8B僅具有22億個激活參數(shù),因此計算成本大大降低——
同時,它還收獲了不錯的性能表現(xiàn)。
如下圖所示:
JetMoE-8B在8個評測基準上獲得了5個sota(包括大模型競技場Open LLM Leaderboard),超過LLaMA-13B、LLaMA2-7B和DeepseekMoE-16B。
在MT-Bench基準上得分6.681,也超過了130億參數(shù)的LLaMA2、Vicuna等模型。
作者介紹
JetMoE一共4位作者,分別是:
- Yikang Shen
MIT-IBM Watson Lab研究員,研究方向NLP。
本碩畢業(yè)于北航,博士經(jīng)歷于Yoshua Bengio創(chuàng)辦的Mila研究機構(gòu)。
- 國振 (Gavin Guo)
MIT博士在讀, 研究方向為3D成像的數(shù)據(jù)高效機器學(xué)習(xí)。
UC伯克利本科畢業(yè),去年夏天作為學(xué)生研究員加入MIT-IBM Watson Lab,導(dǎo)師為Yikang Shen等人。
- 蔡天樂
普林斯頓博士在讀生,本科畢業(yè)于北大應(yīng)用數(shù)學(xué)和計算機科學(xué),目前也是Together.ai 的兼職研究員,與Tri Dao合作。
- Zengyi Qin
MIT博士在讀,同時在創(chuàng)業(yè),MyShell的AI研發(fā)主管。
這家公司剛剛?cè)谫Y了1100萬美元,投資者包括Transformer的作者。
傳送門:https://github.com/myshell-ai/JetMoE
參考鏈接:https://twitter.com/jiayq/status/1775935845205463292