全球最大開源模型再刷爆紀錄!4800億參數(shù)MoE擊敗Llama 3、Mixtral
就在剛剛,擁有128位專家和4800億參數(shù)的Arctic,成功登上了迄今最大開源MoE模型的寶座。
它基于全新的Dense-MoE架構(gòu)設計,由一個10B的稠密Tranformer模型和128×3.66B的MoE MLP組成,并在3.5萬億個token上進行了訓練。
不僅如此,作為一個比「開源」更「開源」的模型,團隊甚至把訓練數(shù)據(jù)的處理方法也全給公開了。
圖片
Arctic的的兩個特點,一個是大,另一個就是非常稀疏。
好處就在于,這種架構(gòu)讓你可以用比別人少好幾倍的訓練開銷,就能得到性能差不多的模型。
也就是說,與其他使用類似計算預算訓練的開源模型相比,Arctic的性能更加優(yōu)異。
比起Llama 3 8B和Llama 2 70B,Arctic所用的訓練計算資源不到它們的一半,評估指標卻取得了相當?shù)姆謹?shù)!
圖片
圖1 編碼(HumanEval+和MBPP+)、SQL生成(Spider) 和指令遵循(IFEval)的企業(yè)智能平均值與訓練成本的比較
具體信息如下——
480B參數(shù),生成期間17B處于活躍狀態(tài);
128位專家,有2位在生成期間活躍;
Instruct & Base版本發(fā)布;
專注于企業(yè)任務(代碼、SQL、推理、跟蹤);
在Apache 2.0下發(fā)布;
FP16精度下約為900GB內(nèi)存,INT4精度下約為240GB內(nèi)存
使用DeepSpeed-MoE訓練。
圖片
主打的就是一個性價比
評測主要看兩項指標,一個是企業(yè)智能指標,一個是學術(shù)基準。
企業(yè)智能指標,是對企業(yè)客戶至關(guān)重要的技能,包括包括編碼(HumanEval+和MBPP+)、SQL生成(Spider)和指令遵循(IFEval)。
同時,團隊也采用了業(yè)界常用的評估LLM的學術(shù)基準,包括世界知識、常識推理和數(shù)學能力。
可以看到,Arctic在多項企業(yè)智能指標中,都超越了Mixtral 8×7B等開源對手。
在計算類別中,它實現(xiàn)了頂級性能,甚至和使用更高計算預算訓練的模型,都有的一拼。
圖片
在學術(shù)基準上,它的表現(xiàn)也不差。
在測評中,團隊發(fā)現(xiàn)了一件有意思的事。
MMLU等世界知識指標,是人們常用的學術(shù)基準測試。而隨著高質(zhì)量網(wǎng)絡和STEM數(shù)據(jù)的增加,MMLU的得分會隨著訓練FLOPS的增加而提高。
但是,Arctic的目標之一,是在保持較小訓練預算的同時優(yōu)化訓練效率,因此,跟其他模型相比,Arctic在MMLU上的得分較低,也實屬情理之中。
由此,如果訓練計算預算高于Arctic的訓練,MMLU性能就將超越Arctic。
當然,MMLU世界知識的性能,并不一定和團隊所關(guān)注的企業(yè)智能直接相關(guān)。
表3 Arctic與DBRX、Llama 3 8B、Llama 3 70B、Mixtral 8x7B、Mixtral 8x22B的對比
企業(yè)級AI的訓練成本,被打下來了!
在以往,用LLM構(gòu)建頂級企業(yè)AI的成本,往往高得離譜,而且需要大量資源,令人望而卻步。
通常,花費的成本高達數(shù)千萬甚至數(shù)億美元,這一成本是驚人的。
如何解決有效訓練和推理的限制?Snowflake AI團隊的研究者一直在做這方面的努力,團隊成員過去曾開源了ZeRO、DeepSpeed、PagedAttention/vLLM和LLM360等系統(tǒng),顯著降低了LLM訓練和推理的成本。
圖片
而今天推出的Arctic,在SQL生成、編碼和遵循基準指令等企業(yè)任務上,表現(xiàn)非常出色。
它為具有成本效益的訓練設定了新的基準,用戶可以以極低的成本,就能創(chuàng)建滿足企業(yè)需求的高質(zhì)量定制模型。
Arctic也是一個真正的開放模型,在Apache 2.0許可下,提供對權(quán)重和代碼的無限制訪問。
從今天開始,Snowflake Arctic就可以從Hugging Face上獲取了。
計算資源僅用一半,表現(xiàn)卻和Llama 3 8B相當
團隊發(fā)現(xiàn),企業(yè)客戶對AI有著一致的需求和使用場景——構(gòu)建對話式SQL數(shù)據(jù)助手、代碼助手和RAG聊天機器人。
為了便于評估,團隊通過對編碼(HumanEval+和MBPP+)、SQL生成(Spider)和指令跟隨(IFEval)取平均值,將這些能力整合到「企業(yè)智能」這個單一指標中。
在開源LLM中,Arctic僅用不到200萬美元(相當于不到3000個GPU周)的訓練計算預算,就實現(xiàn)了頂級的企業(yè)智能。
更重要的是,即使與那些使用顯著更高計算預算訓練的模型相比,它在企業(yè)智能任務上也表現(xiàn)出色。
結(jié)果顯示,Arctic在企業(yè)級評估指標上的表現(xiàn),與Llama 3 8B和Llama 2 70B相當,甚至更優(yōu),而它所使用的訓練計算資源卻不到后兩者的一半。
具體來說,Arctic使用的計算預算只有Llama3 70B的1/17,但在編程(HumanEval+和MBPP+)、SQL(Spider)和指令跟隨(IFEval)等企業(yè)級任務上,都與其不相上下。
表1 Arctic、Llama-2 70B、DBRX和Mixtral 8x22B的模型架構(gòu)和訓練計算量(與活躍參數(shù)和訓練token的乘積成正比)
此外,Arctic的高訓練效率還意味著,Snowflake客戶和整個AI社區(qū)可以以更加經(jīng)濟實惠的方式訓練定制模型。
訓練效率
為了實現(xiàn)如此高的訓練效率,Arctic采用了獨特的Dense-MoE Hybrid transformer架構(gòu)。
該架構(gòu)將一個10B規(guī)模的稠密Transformer模型與一個128×3.66B規(guī)模的殘差MoE MLP相結(jié)合,雖然總參數(shù)量達到480B,但通過top-2 gating的方式只選擇了其中17B個參數(shù)保持活躍。
Arctic的設計和訓練基于以下三個關(guān)鍵創(chuàng)新:
1. 更多但精煉的專家,以及更多的專家選擇
首先,DeepSpeed團隊在2021年末便證明了,MoE(Mixture of Experts)可以在不增加計算成本的情況下,顯著提高LLM模型的質(zhì)量。
其次,模型質(zhì)量的提升主要取決于MoE模型中專家的數(shù)量、總參數(shù)量以及這些專家可以組合在一起的方式和數(shù)量。
基于此,Arctic被設計為擁有480B個參數(shù),分布在128個細粒度專家中,并使用top-2 gating選擇17B個活躍參數(shù)。相比之下,最近的MoE模型使用的專家數(shù)量就要少得多了(如表2所示)。
從直觀上看,Arctic利用更大的總參數(shù)量和眾多專家來擴大模型容量,同時更明智地在眾多精煉的專家中進行選擇,并使用適度數(shù)量的活躍參數(shù)來實現(xiàn)資源高效的訓練和推理,最終獲得頂級的智能。
圖2 標準MoE架構(gòu) vs. Arctic
2. 架構(gòu)和系統(tǒng)協(xié)同設計
即便是用最強大的AI硬件,想要基于普通的MoE架構(gòu)訓練大量專家效率依然很低。
其原因在于,專家之間存在的全通信開銷非常高昂。不過,如果能將通信與計算重疊,那么就可以極大地降低這種開銷。
因此,團隊在Arctic架構(gòu)中將一個密集的Transformer與一個殘差MoE組件(圖2)相結(jié)合,從而使系統(tǒng)能夠通過通信計算重疊來消除大部分通信開銷,最終實現(xiàn)了極佳的訓練效率。
3. 面向企業(yè)的數(shù)據(jù)課程
要在代碼生成和SQL等企業(yè)指標上表現(xiàn)出色,需要與訓練通用指標的模型截然不同的數(shù)據(jù)課程。
團隊在進行了數(shù)百次小規(guī)模的對比實驗后發(fā)現(xiàn),常識推理等通用技能可以在開始時學習,而編碼、數(shù)學和SQL等更復雜的指標可以在訓練的后期有效學習。
因此,Arctic采用了三階段課程進行訓練,每個階段的數(shù)據(jù)組成不同——
第一階段(1T Tokens)側(cè)重于通用技能,后兩個階段(1.5T和1T Tokens)側(cè)重于企業(yè)級技能。
表2 Arctic三階段訓練的動態(tài)數(shù)據(jù)組成
推理效率
訓練效率,只是Arctic高效的其中一個方面。
如果希望低成本部署模型,推理效率也同樣至關(guān)重要。
作為MoE模型規(guī)模的飛躍,Arctic使用了比其他開源自回歸模型更多的專家和參數(shù)。
因此,為了有效地在Arctic上運行推理,團隊做了一些系統(tǒng)性的創(chuàng)新——
a) 在較小batch的交互式推理中(比如批大小為1),MoE模型的推理延遲受到了讀取所有活躍參數(shù)所需時間的瓶頸,其中,推理是受內(nèi)存帶寬限制的。
在這樣的批大小下,Arctic(17B活躍參數(shù))的內(nèi)存讀取次數(shù)比Code-Llama 70B少4倍,比 Mixtral 8x22B(44B活動參數(shù))少2.5倍,從而實現(xiàn)更快的推理性能。
為此,團隊跟英偉達的TensorRT-LLM和vLLM團隊展開合作,為交互式推理提供了Arctic的初步實現(xiàn)。
通過FP8量化,團隊可以將Arctic放入單個GPU節(jié)點中。
雖然仍遠未完全優(yōu)化,但在批大小為1時,Arctic的吞吐量超過70+token/秒,這樣就實現(xiàn)了有效的交互式服務。
b) 當批大小的規(guī)模顯著增加,例如每次前向傳遞要處理數(shù)千個token時,Arctic就會從內(nèi)存帶寬受限轉(zhuǎn)變?yōu)橛嬎闶芟蓿藭r推理的瓶頸就在于每個token的活躍參數(shù)。
在這一點上,與CodeLlama 70B和Llama 3 70B相比,Arctic的計算需求減少了4倍。
為了實現(xiàn)計算受限的推理和與Arctic中活躍參數(shù)數(shù)量較少相對應的高吞吐量(如下圖所示),需要較大的batch size。
要實現(xiàn)這一點,需要有足夠的KV緩存內(nèi)存來支持較大的batch size,同時也需要足夠的內(nèi)存來存儲近500B的模型參數(shù)。
面對這重重挑戰(zhàn),最終團隊還是找到了辦法。
通過使用FP8權(quán)重、分割融合和連續(xù)批處理、節(jié)點內(nèi)的張量并行性以及節(jié)點間的管線并行性等系統(tǒng)優(yōu)化組合,團隊在雙節(jié)點推理中,實現(xiàn)了這一目標。
圖3 推理期間編碼(HumanEval+和MBPP+)、SQL生成(Spider)和指令跟蹤 (IFEval)企業(yè)智能的平均值與活躍參數(shù)的對比
開源代碼
新模型Arctic基礎(chǔ)模型和指令微調(diào)模型代碼全部開源,任何人可以將其用于研究、產(chǎn)品、原型當中。
項目地址:https://github.com/Snowflake-Labs/snowflake-arctic
研究人員基于LoRA的微調(diào)的pipeline和配方(recipe),并允許在單個節(jié)點上進行高效的模型微調(diào)。
現(xiàn)在,Snowflake正在與英偉達TensorRT-LLM和vLLM開展合作,為Arctic模型開發(fā)初始的推理實現(xiàn),并且針對批大小為1的交互式使用進行了優(yōu)化。
未來,他們還將與社區(qū)合作,解決真正大型MoE更大的批大小的推理復雜性。
Cookbook:https://medium.com/snowflake/snowflake-arctic-cookbook-series-exploring-mixture-of-experts-moe-c7d6b8f14d16
另外,Arctic現(xiàn)使用的是4k上下文窗口進行訓練,研究人員還將開發(fā)一種基于注意力下沉(attention-sinks)的滑動窗口的方法,以支持未來幾周無限序列生成能力。
下一步,將會擴展到32K上下文窗口。
團隊介紹
Snowflake的CEO,是Sridhar Ramaswamy,是前谷歌高級副總裁。
圖片
在谷歌工作15年后,他成為Neeva的聯(lián)合創(chuàng)始人,后來Neeva被Snowflake收購。
圖片
圖片
圖片
他在印度理工學院馬德拉斯分校獲得計算機學士學位,并在布朗大學獲得計算機博士學位。
圖片
AI團隊的一把手Vivek Raghunathan,也是前谷歌副總裁。
圖片
他曾擔任微軟研究員,后在谷歌從事機器學習、廣告基礎(chǔ)架構(gòu)等方面工作,18年開始在谷歌擔任副總裁,領(lǐng)導YouTube團隊。
隨后,他和Sridhar Ramaswamy共同創(chuàng)辦了Neeva。
圖片
圖片
Raghunathan同樣也是印度理工學院的校友,不過是在孟買分校獲得的學士學位。之后,他在UIUC取得了碩士和博士學位。
圖片
為了發(fā)展AI,兩人把DeepSpeed團隊最頂尖的幾個元老都挖了過來,包括Zhewei Yao和Yuxiong He。
Zhewei Yao在UC伯克利獲得博士學位,研究興趣在于計算統(tǒng)計、優(yōu)化和機器學習。(在此之前,他2016年曾獲得上交大數(shù)學學士學位。)
圖片
他從2021年開始便加入了微軟,在微軟擔任首席研究員和研發(fā)經(jīng)理,致力于高效的大規(guī)模訓練和推理。
目前,他是Snowflake的高級科學家和SDE II,同時也是Snowflake大規(guī)模預訓練創(chuàng)始成員。
圖片
Yuxiong He在微軟任職13年,是DeepSpeed的創(chuàng)始人之一,最近加入了Snowflake。
她曾在新加坡南陽理工大學獲得了計算機工程學士學位。
圖片
團隊的另一位華人大牛Aurick Qiao,去年11月剛加入Snowflake。
圖片
CMU讀博期間,他曾獲得Osdi 2022的最佳論文優(yōu)勝獎。此前曾在微軟、Dropbox工作。
曾擔任Petuum CEO,以及LMNet的聯(lián)合創(chuàng)始人。
圖片
Hao Zhang是UCSD的Hal?c?o?數(shù)據(jù)科學研究所和計算機科學與工程系的助理教授。
圖片
他曾獲得了CMU計算機博士學位,師從Eric Xing。在攻讀博士學位期間,他休學一段時間并在ML平臺初創(chuàng)公司Petuum工作。
Hao Zhang在2023年聯(lián)合創(chuàng)立了LMnet.ai,這家公司于同年11月加入了Snowflake。
他此前還共同創(chuàng)辦了非營利組織LMSYS Org,該組織訓練了此前十分流行的模型Vicuna以及發(fā)起和維護目前最重要的大語言模型評測機制:Chatbot Arena。
他本人的研究興趣是機器學習與系統(tǒng)的交叉領(lǐng)域。
參考資料:https://www.snowflake.com/blog/arctic-open-efficient-foundation-language-models-snowflake/