Attention isn’t all you need!Mamba混合大模型開源:三倍Transformer吞吐量

机器之心發表於2024-03-29

Mamba 時代來了?

自 2017 年開創性研究論文《Attention is All You Need》問世以來,transformer 架構就一直主導著生成式人工智慧領域。

然而,transformer 架構實際上有兩個顯著缺點:

  • 記憶體佔用大:Transformer 的記憶體佔用量隨上下文長度而變化。這使得在沒有大量硬體資源的情況下執行長上下文視窗或大量並行批處理變得具有挑戰性,從而限制了廣泛的實驗和部署。

  • 隨著上下文長度的增加,推理速度會變慢:Transformer 的注意力機制隨序列長度呈二次方擴充套件,並且會降低吞吐量,因為每個 token 都依賴於它之前的整個序列,從而將長上下文用例置於高效生產的範圍之外。

但 transformer 並不是生成式人工智慧唯一的前進方向。最近,AI21 Labs 推出並開源了一種名為「Jamba」的新方法,在多個基準上超越了 transformer。

圖片

Hugging Face 地址:https://huggingface.co/ai21labs/Jamba-v0.1

圖片

Mamba 的 SSM 架構可以很好地解決 transformer 的記憶體資源和上下文問題。然而,Mamba 方法很難提供與 transformer 模型相同的輸出水平。

Jamba 將基於結構化狀態空間模型 (SSM) 的 Mamba 模型與 transformer 架構相結合,旨在將 SSM 和 transformer 的最佳屬性結合在一起。

圖片

Jamba 還可以作為 NVIDIA NIM 推理微服務從 NVIDIA API 目錄進行訪問,企業應用程式開發人員可以使用 NVIDIA AI Enterprise 軟體平臺進行部署。

總的來說,Jamba 模型具有以下特點:

  • 第一個基於 Mamba 的生產級模型,採用新穎的 SSM-Transformer 混合架構;

  • 與 Mixtral 8x7B 相比,長上下文上的吞吐量提高了 3 倍;

  • 提供對 256K 上下文視窗的訪問;

  • 公開了模型權重

  • 同等引數規模中唯一能夠在單個 GPU 上容納高達 140K 上下文的模型。

模型架構

如下圖所示,Jamba 的架構採用塊層(blocks-and-layers)方法,使 Jamba 能夠整合兩種架構。每個 Jamba 塊包含一個注意力層或一個 Mamba 層,後跟一個多層感知器(MLP),從而形成 transformer 層。

圖片

Jamba 利用 MoE 來增加模型引數的總數,同時簡化推理中使用的活躍引數的數量,從而在計算需求沒有相應增加的情況下獲得更高的模型容量。為了在單個 80GB GPU 上最大限度地提高模型的質量和吞吐量,研究團隊最佳化了所使用的 MoE 層和專家的數量,為常見推理工作負載留出了足夠的記憶體。

Jamba 的 MoE 層允許它在推理時僅利用可用的 52B 引數中的 12B,並且其混合架構使這些 12B 活躍引數比同等大小的純 transformer 模型更有效。

此前,沒有人將 Mamba 擴充套件到 3B 引數之外。Jamba 是同類模型中第一個達到生產級規模的混合架構。

吞吐量和效率

初步評估實驗表明,Jamba 在吞吐量和效率等關鍵衡量指標上表現出色。

在效率方面,Jamba 在長上下文上的吞吐量達到了 Mixtral 8x7B 的 3 倍。Jamba 比 Mixtral 8x7B 等大小相當的基於 Transformer 的模型更高效。

圖片

在成本方面,Jamba 可以在單個 GPU 上容納 140K 上下文。與當前類似大小的其他開源模型相比,Jamba 能提供更多的部署和實驗機會。

圖片

需要注意的是,Jamba 目前不太可能取代當前基於 Transformer 的大型語言模型 (LLM),但它可能會成為某些領域的補充。

參考連結:

https://www.ai21.com/blog/announcing-jamba

https://venturebeat.com/ai/ai21-labs-juices-up-gen-ai-transformers-with-jamba/

相關文章