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