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