【新智元導讀】最大開源模型,再次刷爆紀錄!Snowflake 的 Arctic,以 128 位專家和 4800 億參數,成爲迄今最大的開源模型。它的特點,是又大又稀疏,因此計算資源隻用了不到 Llama 3 8B 的一半,就達到了相同的性能指标。
就在剛剛,擁有 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)的企業智能平均值與訓練成本的比較
具體信息如下——
480B 參數,生成期間 17B 處于活躍狀态;
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 三階段訓練的動态數據組成
推理效率
訓練效率,隻是 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。
他本人的研究興趣是機器學習與系統的交叉領域。
參考資料:
https://www.snowflake.com/blog/arctic-open-efficient-foundation-language-models-snowflake/