精彩精彩，第一個把爆火 Mamba 架構真正擴展到足夠大的工作來了。
520 億參數，還是Mamba+Transformer混合架構。
它的名字叫Jamba。
取兩種架構之長，模型質量和效率兼得，要吞吐量有吞吐量，要低内存有低内存。
初步跑分顯示：
Jamba 性能總體接近Mixtral 8x-7B，處理 128k 長上下文時吞吐量卻是其 3 倍。
一共支持 256k 上下文，而單張 A100 GPU 即可處理 140k，直接拿下同等規模模型之最，高效又經濟。
這項成果，來自以色列 AI 公司AI21labs。
Mamba 原作者看了之後都激動轉發：
絕對的 " 大新聞 "。
Mamba、Transformer，合體
由 CMU 和普林斯頓大學提出的 Mamba，解決了 Transformer 的局限性（随着推理上下文越長，模型内存占用量越大，同時推理速度變慢，由此導緻算力消耗巨大）。
但它也有自己的缺點——
在不關注整個上下文的情況下，Mamba 的輸出質量很差，尤其是在召回相關的任務上。
本着 " 既要也要 " 的原則，Jamba 站出來提供兩全其美之作。
Jamba 由 Transformer、Mamba 和 MoE 層組成，可同時優化内存、吞吐量和性能。
如下圖所示，爲了集成兩種架構，Jamba 采用塊層（blocks-and-layers）組合的創新方法。
簡單來說，就是每個 Jamba 塊包含一個注意力層或一個 Mamba 層，再跟一個多層感知器 MLP，總體比例保證爲每八層一個 Transformer 層。
其次，Jamba 利用 MoE 來增加模型參數的總量，同時簡化推理中使用的活動參數量。
最終模型容量高了，計算需求也沒有相應的增加。
而爲了在單張 GPU（80GB）上最大限度地提高模型吞吐量，Jamba 還優化了所用 MoE 層和專家數量，最終爲日常推理工作負載留出足夠内存。
值得一提的是，在推理時，Jamba 的 MoE 層僅需 520 億可用參數中的 120 億，就能同時保證比同等大小的僅 Transformer 模型更高效。
要知道，此前有人光是嘗試過擴展 Mamba，就沒能做到 30 億參數之上。
因此，除了成功合體 Mamba 和 Transformer，Jamba 也達成了第二大成就：
同類中第一個達到生産級規模和質量的混合架構（SSM 混 Transformer）（ps. Mamba 就是一種狀态空間模型 SSM）。
吞吐量和效率 up
初步評估顯示，Jamba 在吞吐量和效率等關鍵指标上表現出色。
首先，Jamba 可以在長上下文中提供 3 倍吞吐量，比 Mixtral 8x7B 等大小相當的 Transformer 模型都要高效。
如下圖所示，當上下文窗口達到 128k 時，Jamba 的每秒 token 數近乎 1500，而此時表現最好的 Mixtral 8x7B 應該才在 500 往上的樣子。
其次，在單張 GPU 上，Jamba 最多可以容納 140k 上下文，經濟又高效。
相比之下，Mixtral 8x7B 爲 64k，Llama2 70B 則僅爲 16k。
第三，Jamba 的輸出質量也得到了保證。
在如下一系列推理基準上，4 項中有 3 項它都拿下了 SOTA。同時，在 GSM8K 等基準上，Jamba 即使沒有奪魁，也和 SOTA 模型打了個不相上下。
總體來說，Jamba 的性能接近 Mixtral 8x7B。
最後，作者提示，别忘了，這些都還隻是初步改造後的結果，後續還有很多優化空間（比如 MoE 并行、更快的 Mamba 實現）。所以到時性能會更強。
好消息：Jamba 現在已經上線 Hugging Face，并且劃重點：采用 apache-2.0 許可。
（Jamba 的指令版本則将很快通過 AI21labs 平台上線。）
網友看完都感動哭了。
傳送門：
https://huggingface.co/ai21labs/Jamba-v0.1
參考鏈接：
[ 1 ] https://www.ai21.com/blog/announcing-jamba
[ 2 ] https://www.ai21.com/jamba
[ 3 ] https://twitter.com/AI21Labs/status/1773350888427438424?s=20
[ 4 ] https://twitter.com/tri_dao/status/1773418926518734957?s=20
