Jamba前生今世：1.5開源來襲
AI21服務(wù)于企業(yè)，為企業(yè)構(gòu)建基礎(chǔ)模型和AI系統(tǒng)以加速GenAI在生產(chǎn)中的使用。AI21 成立于2017年，已從NVIDIA、Intel、Google等公司共籌集了3.36億美元。它是最早將生成式AI推向大眾的公司之一，借助AI21平臺，企業(yè)可以構(gòu)建自己的生成式AI的應(yīng)用程序和服務(wù)。
Jamba算是世界上第一款基于Mamba的生產(chǎn)級別的模型。近日AI21 Labs發(fā)布Jamba 1.5系列開放模型，包括Jamba 1.5 Mini和Jamba 1.5 Large，這些模型建立在新穎的SSM-Transformer架構(gòu)之上，某種意義上也算是AI技術(shù)的突破。
1.Jamba
Jamba 1.5創(chuàng)新的核心是其令人印象深刻的256K上下文窗口，這是目前最長上下文的開放模型。在單個實例中處理和理解如此大量信息的能力允許在各種任務(wù)中實現(xiàn)更細致、更全面和更準確的輸出。Jamba 1.5 Mini可以在單個GPU上處理高達140K令牌的上下文長度。
圖片
AI21 Labs的原話是：“大多數(shù)其他模型聲稱具有較長的上下文窗口，但無法在其上下文窗口的上限處保持相同的性能質(zhì)量，而Jamba 1.5系列在其 256K上下文窗口的整個跨度內(nèi)保持其長上下文處理。”
能夠有效處理很長的上下文模型對于幾乎企業(yè)級生成式的應(yīng)用程序都至關(guān)重要。除了徹底、精確地總結(jié)和分析冗長的文檔外，長上下文模型還通過消除連續(xù)分塊和重復檢索的需要，大大提高了RAG和代理工作流程的質(zhì)量，并降低了它們的成本。
圖片
上圖為Jamba 1.5 Mini在2片A100 80GB GPU進行延遲評估， 而Jamba 1.5 Large則是在8塊A100 80GB GPU上完成的，整體表現(xiàn)不錯。下圖為
圖片
上表是Jamba和其他模型整體參數(shù)規(guī)模和實際運行激活參數(shù)規(guī)模的對比。Jamba 1.5 Large算是巔峰之作，它雖然具備398B的總參數(shù)，然而運行時僅僅有94B的活躍參數(shù)。而Jamba 1.5 Mini具有52B的總參數(shù)和12B的活躍參數(shù)。<MoE!!>
2.SSM-Transformer
Jamba 1.5基于混合SSM-Transformer構(gòu)建，它結(jié)合了兩種強大的AI 范式的優(yōu)勢：Mamba和Transformer。讀者是否還記得Mamba系列？！Mamba為模型帶來了卓越的效率和速度，而Transformer架構(gòu)則在各種 NLP任務(wù)中保持高質(zhì)量輸出和穩(wěn)健的性能。
圖片
在Jamba的這個版本中，AI21試驗了Mamba-2，這是Mamba的更快和改進版本，它的性能分別優(yōu)于Mamba和Transformers。在實驗中如上圖所示在混合架構(gòu)中，Mamba-1+Attention組合比Mamba-2+Attention效果更好，因此本次在Jamba-1.5-Large中使用Mamba-1。
當然在這個過程中還發(fā)現(xiàn)混合架構(gòu)的性能優(yōu)于純Mamba-2。通過推測可能是因為在 Mamba層之間交錯有完整的Tranformer層時，Mamba-2相對于Mamba-1 的優(yōu)勢就不那么明顯了，Mamba-2雖然可以具備更大狀態(tài)存儲能力，然而Mamba-1通過Transformer可以匯集來自整個上下文的信息。
圖片
上圖為最早期的Jamba論文中Jamba Block的架構(gòu)圖。它為八層的架構(gòu)，Mamba層和Transformer層的比例為7:1，每兩層運用了MoE。
而Jamba-1.5-Large遵循相同的Jamba結(jié)構(gòu)，但容量更大。它具有94B活動參數(shù)和398B總參數(shù)。它有9個區(qū)塊，每個區(qū)塊都有以下規(guī)格：
- l=8，每塊具有8層
- a:m=1:7，這個比例是研究人員在對Jamba研究中發(fā)現(xiàn)的最佳的比例
- e=2，每兩層都使用MoE代替單個MLP
- n=16，有16個專家，在每個Token上選擇頭部的兩個，即K=2。
- 隱藏維度為8192
- Q為64頭，KV為8頭
3.ExpertsInt8
Jamba 1.5 Large 背后的關(guān)鍵創(chuàng)新之一是引入ExpertsInt8，這是一種專為專家混合 （MoE） 模型設(shè)計的新型量化技術(shù)。這種技術(shù)可以在A100 GPU 上使用，而FP8僅在H100上可以使用。當然這種技術(shù)最重要的特點如下：
- 快速的量化過程，只需幾分鐘，而不是幾小時或幾天
- 不依賴校準，增強了穩(wěn)定性和易用性
- 能夠?qū)?BF16 用于大型激活，在最重要的地方保持精度
- 在 MoE 模型的 vLLM 量化技術(shù)中延遲最低，且不會犧牲質(zhì)量
因為超過85%的模型權(quán)重在MoE層中，超過90%在 MoE或MLP層中，那么如何量化這些權(quán)重且同時仍然享受快速BF16內(nèi)核的好處？研究人員將 MoE和MLP權(quán)重量化為INT8，將它們保存在INT8中，并在實際計算之前將它們返回BF16。
最重要的是，逆量化步驟直接發(fā)生在vLLM的 fused_moe內(nèi)核內(nèi)部。通過這種方式，逆量化過程增加的開銷可以忽略不計，甚至導致比BF16更高的延遲。這個修改的fused_moe內(nèi)核目前已經(jīng)貢獻給vLLM。
圖片
上圖為不同量化技術(shù)比較，顯示1024個令牌上下文和128個令牌解碼的端到端延遲。ExpertsInt8的性能與FP8相似，但應(yīng)用起來快速而簡單，并且依舊享受BF16激活，最重要的是能夠適用于FP8不可用的A100 GPU。
圖片
至于在一些通用基準的評測上，Jamba-1.5也獲得不俗的表現(xiàn)。與同級別的模型相比，性能相似然而擁有更好的吞吐量和低延遲的能力。
