僅需Llama3 1/17的訓練成本，Snowflake開源128x3B MoE模型
机器之心發表於2024-04-25
Snowflake 釋出高「企業智慧」模型 Arctic，專注於企業內部應用。剛剛，資料管理和倉庫提供商 Snowflake 宣佈加入 LLM 混戰，釋出了一款專注於企業級應用的頂級大型語言模型（LLM）——Snowflake Arctic。作為一家雲端計算公司推出的 LLM，Arctic 主要具備以下兩個方面的優勢：- 高效智慧：Arctic 在企業任務方面表現出色，例如 SQL 生成、程式設計和指令遵循，甚至可與使用更高計算成本訓練的開源模型媲美。Arctic 為經濟高效的訓練設定了新的基線，使 Snowflake 客戶能夠以低成本為其企業需求建立高質量的定製模型。
- 開源開放：Arctic 採用 Apache 2.0 許可，提供對權重和程式碼的開放訪問，Snowflake 還將開源所有的資料方案和研究發現。
現在，你可以在 Hugging Face 上訪問 Arctic 模型。Snowflake 表示：使用者很快可以透過一些模型庫獲取，包括 Snowflake Cortex、AWS、微軟 Azure、NVIDIA API、Lamini、Perplexity、Replicate 和 Together 等。Hugging Face：https://huggingface.co/Snowflake/snowflake-arctic-instructArctic 的上下文視窗設定為 4K，研究團隊正在研發基於注意力池（attention-sink）的滑動視窗實現，在未來幾周內將支援無限序列生成，並在不久的將來擴充套件到 32K 注意力視窗。Snowflake 的研究團隊從企業客戶的 AI 需求和使用案例中看到了一個一致的模式：企業希望使用 LLM 構建對話式 SQL 資料 copilot、程式碼 copilot 和 RAG 聊天機器人。這意味著 LLM 需要在 SQL、程式碼、複雜指令遵循和生成具體響應方面表現出色。Snowflake 將這些能力融合成一個稱為「企業智慧」的單一指標，具體方式是對編碼（HumanEval + 和 MBPP+）、SQL 生成（Spider）和指令遵循（IFEval）效能水平取平均值。Arctic 在開源 LLM 中達到了頂級的「企業智慧」水平，而且是在大約不到 200 萬美元的訓練計算成本（少於 3K GPU 周）的情況下做到的。這意味著 Arctic 比其他使用類似計算成本訓練的開源模型能力更強。更重要的是，即使與那些使用遠高於其的計算成本訓練的模型相比，Arctic 在企業智慧方面也表現出色。Arctic 的高訓練效率意味著 Snowflake 的客戶和整個 AI 社群可以以更經濟的方式訓練定製模型。如圖 1 所示，Arctic 在企業智慧指標上與 LLAMA 3 8B 和 LLAMA 2 70B 不相上下，而使用的訓練計算成本不到一半。並且，儘管僅使用 1/17 倍的計算成本，Arctic 在編碼（HumanEval + 和 MBPP+）、SQL（Spider）和指令遵循（IFEval）等指標上可與 Llama3 70B 媲美，即 Arctic 在保持整體效能競爭力的同時做到了這一點。此外，Snowflake 還在學術基準上評估了 Arctic，涉及世界知識、常識推理和數學能力，完整評估結果如下圖所示：為了達到上述訓練效率，Arctic 採用一種獨特的 Dense-MoE 混合 transformer 架構。它將一個 10B 的密集 transformer 模型與一個 128×3.66B 的殘差 MoE MLP 結合起來，總共有 480B 引數和 17B 活躍引數，使用 top-2 gating 來進行選擇。設計和訓練 Arctic 時，研究團隊使用了以下三個關鍵的見解和創新：2021 年底，DeepSpeed 團隊證明了 MoE 可以應用於自迴歸 LLM，從而顯著提高模型質量而不增加計算成本。在設計 Arctic 時，研究團隊注意到，基於這個思路，模型質量的提高主要取決於 MoE 模型中的專家數量和總引數量，以及這些專家的組合方式數量。基於此，Arctic 被設計為在 128 個細粒度（fine-grained）專家之間分佈 480B 引數，並使用 top-2 gating 來選擇 17B 活躍引數。在強大的 AI 訓練硬體上訓練具有大量專家的基本 MoE 架構非常低效，因為專家之間的全連線通訊開銷很高。Snowflake 發現，如果通訊可以與計算重疊，就可以省去這種開銷。因此，Arctic 將密集 transformer 與殘差 MoE 元件相結合（圖 2），透過通訊計算重疊，使訓練系統能夠實現良好的訓練效率，隱藏了通訊開銷的大部分。在程式碼生成和 SQL 等企業級指標上表現出色需要與通用指標截然不同的資料課程學習（Curriculum Learning）。透過數百次小規模的消融實驗，該團隊瞭解到通用技能，如常識推理，可以在初始階段學習；而編碼、數學和 SQL 等更復雜的指標可以在訓練後期有效學習。這可以類比於人類的生活教育，從簡單到困難逐步獲取能力。因此，Arctic 使用一個三階段的課程學習，每個階段的資料構成都不同，第一階段側重於通用技能（1T token），後兩個階段側重於企業技能（1.5T 和 1T token）。推理效率也是模型高效的一個重要方面，影響到模型是否可以在低成本下進行實際部署。Arctic 代表了 MoE 模型規模的一次飛躍，它比任何其他開源自迴歸 MoE 模型都使用了更多的專家和總引數。因此，Snowflake 需要幾個創新思路來確保 Arctic 能夠高效推理：a) 在批大小較小的互動推理中，例如批大小為 1，MoE 模型的推理延遲受制於讀取所有活躍引數的時間，推理是受記憶體頻寬限制的。在這種批大小下，Arctic（17B 活躍引數）的記憶體讀取量僅為 Code-Llama 70B 的 1/4、Mixtral 8x22B（44B 活躍引數）的 2/5，從而具備更快的推理速率。b) 當批大小顯著增加，例如每次前向傳遞數千個 token 時，Arctic 從記憶體頻寬受限轉變為計算受限，推理受到每個 token 的活躍引數的限制。在這方面，Arctic 的計算量是 CodeLlama 70B 和 Llama 3 70B 的 1/4。為了實現計算受限的推理和與 Arctic 中少量活躍引數相匹配的高吞吐量，需要一個較大的批大小。實現這一點需要有足夠的 KV 快取來支援，同時還需要足夠的記憶體來儲存模型的近 500B 引數。雖然具有挑戰性，但 Snowflake 透過使用兩個節點進行推理，並結合 FP8 權重、split-fuse 和連續批處理、節點內張量並行以及節點間 pipeline 並行等系統最佳化來實現。研究團隊已與 NVIDIA 展開密切合作，針對由 TensorRT-LLM 驅動的 NVIDIA NIM 微服務進行推理最佳化。同時，研究團隊還與 vLLM 社群合作，內部開發團隊也將在未來幾周內為企業用例實現 Arctic 的高效推理。參考連結：https://www.snowflake.com/blog/arctic-open-efficient-foundation-language-models-snowflake/