Jamba-1.5：大規(guī)?；旌蟃ransformer-Mamba模型
一、結論寫在前面
論文標題：Jamba-1.5: Hybrid Transformer-Mamba Models at Scale
論文鏈接：https://arxiv.org/pdf/2408.12570
模型：https://huggingface.co/ai21labs
論文介紹了Jamba-1.5，基于Jamba架構的新型指令調優(yōu)大型語言模型。Jamba是一種混合Transformer-Mamba專家混合架構，能夠在不同上下文長度下提供高吞吐量和低內存使用，同時保持與Transformer模型相同或更好的質量。
論文發(fā)布了兩種模型尺寸：Jamba-1.5-Large，具有940億活躍參數(shù)，以及Jamba-1.5-Mini，具有120億活躍參數(shù)。這兩種模型都針對多種對話和指令遵循能力進行了微調，并且具有256Ktoken的有效上下文長度，是開放權重模型中最大的。
為了支持成本效益高的推理，論文引入了ExpertsInt8，一種新穎的量化技術，允許在處理256K token上下文時，將Jamba-1.5-Large適配到具有8張80GB GPU的機器上，而不損失質量。在學術和聊天機器人基準測試中評估時，Jamba模型取得了優(yōu)異的成績，同時提供了高吞吐量，并在長上下文基準測試中超越了其他開放權重模型。
二、論文的簡單介紹
2.1 論文的背景?
論文介紹了Jamba-1.5，兩個基于論文的Jamba架構[的新型大語言模型，可供公眾使用。Jamba-1.5-Mini是論文早期Jamba版本的更新和指令調優(yōu)版本。與其較小的同類產(chǎn)品一樣，Jamba-1.5-Large是一種混合架構，結合了Transformer和Mamba層，以及專家混合(MoE)模塊。
自Jamba推出以來，類似的努力已經(jīng)證實了在高達8B參數(shù)規(guī)模上結合Transformer和狀態(tài)空間模型的優(yōu)勢。Jamba-1.5-Large在更大規(guī)模上展示了這種架構的優(yōu)勢。它具有94B活躍參數(shù)，總共398B參數(shù)。即使在這個大尺寸下，由于Jamba架構的效率以及論文開發(fā)的一種新穎量化技術，該模型可以在處理256Ktoken上下文時適配到具有8張80GB GPU的單臺機器上。
Jamba-1.5-Mini和Jamba-1.5-Large均為經(jīng)過指令微調的模型，通過Post-training賦予了多種能力。論文在廣泛的基準測試中評估發(fā)現(xiàn)，它們在性能上與同尺寸模型相當，同時得益于Jamba架構的高效性。特別是，Jamba-1.5模型在長上下文評估中表現(xiàn)突出，使其成為唯一在RULER
基準測試中有效長度達到256K的模型，同時實現(xiàn)了KV緩存內存減少10倍以及卓越的吞吐量和延遲。
這些模型已公開可用：
Jamba-1.5-Mini: https://huggingface.co/ai21labs/AI21-Jamba-1.5-Mini
Jamba-1.5-Large: https://huggingface.co/ai21labs/AI21-Jamba-1.5-Large
2.2 模型架構
Jamba-1.5-Large基于Jamba ，這是一種混合解碼器架構，融合了Transformer層 與Mamba層（一種狀態(tài)空間模型（state-space model，SSM）），并附加了混合專家（MoE）模塊 。詳見 [24] 對該架構的詳細描述。
在開發(fā)Jamba [24] 的過程中，論文發(fā)現(xiàn)Transformer、Mamba和MoE元素的結合有助于平衡吞吐量、內存使用和質量的需求。Jamba-1.5-Large在大規(guī)模上展現(xiàn)了這種靈活性。
Jamba-1.5-Large 遵循相同的 Jamba 結構，但容量更大。它擁有 94B 活躍參數(shù)和 398B 總參數(shù)。它包含 9 個塊，每個塊具有以下規(guī)格：
?l= 8每個block包含 8 層。
?a : m=1 : 7注意力層與 Mamba 層的比例。在論文的 Jamba 研究中，這一比例被發(fā)現(xiàn)是最佳的 [ 2 4 ]，后續(xù)工作 [6, 37] 也證實了類似比例的成功.
?每隔 e=2 層使用 MoE 替代單一 MLP。共有 n=1 6 個專家，每個token選擇前 K=2 個。
?隱藏狀態(tài)維度為 8192。
?注意力查詢頭數(shù)為 64，KV 頭數(shù)為 8。
表 1 將 Jamba-1.5 模型與公開可用的類似尺寸模型進行了比較。Jamba-1.5-Mini 的活躍參數(shù)數(shù)量與 Mixtral 8x7B 相近，而 Jamba-1.5-Large 的活躍參數(shù)數(shù)量介于 LLaMA-3.1-70B 和 Mistral-Large-2 之間。同時，論文的兩個 Jamba 模型在 KV 緩存內存使用（256K token）方面遠小于所有其他模型，相較于各自的對應模型，大約減少了近一個數(shù)量級的內存使用。
通過這些設置以及論文的專用量化（第 3.1 節(jié)），Jamba-1.5-Large 可以在單臺配備 8 塊 80GB GPU 的機器上提供服務，上下文長度可達 256K token。
表 1：Jamba-1.5-Mini、Jamba-1.5-Large 與近期開放模型在總可用參數(shù)、活躍參數(shù)及長上下文 KV 緩存內存方面的比較。Jamba-1.5-Mini 和 Jamba-1.5-Large 在 KV 緩存內存需求上提供了顯著的減少。
對于這次發(fā)布，論文還嘗試了Mamba-2 [6]，這是Mamba的一個更快且改進的版本，據(jù)報道其性能優(yōu)于單獨的Mamba和Transformers。然而，如圖1所示，論文發(fā)現(xiàn)，在混合架構中，Mamba-1-Attention組合的性能優(yōu)于Mamba-2-Attention，因此論文在Jamba-1.5-Large中采用了Mamba-1。（論文還發(fā)現(xiàn)混合架構的性能優(yōu)于純Mamba-2。）論文推測這是因為Mamba-2相對于Mamba-1的一些優(yōu)勢，特別是能夠使用更大的狀態(tài)大小，在論文將全注意力層交錯放置在Mamba層之間時，其重要性有所降低，因為這些全注意力層能夠從整個上下文中匯聚信息。
圖1：Mamba-1、Mamba-2、Mamba-1-Attention和Mamba-2-Attention在訓練100B tokens模型上的比較。盡管Mamba-2在沒有注意力機制的情況下優(yōu)于Mamba-1，但混合的Mamba-1-Attention表現(xiàn)更佳。
2.3 服務考慮與改進
論文分享了一些見解和改進措施，以實現(xiàn)Jamba模型在大規(guī)模上的高效服務。
2.3.1 專家Int8量化
為了支持Jamba-1.5-Large的高效服務，論文開發(fā)了一種新的量化技術，論文稱之為ExpertsInt8。論文觀察到，超過85%的模型權重位于MoE層中，超過90%位于MoE或MLP層中。論文希望量化這些權重，同時仍然享受快速BF16內核的好處。為此，論文將MoE和MLP權重量化為INT8，以INT8格式保存，并在實際計算前將其反量化回BF16。重要的是，反量化步驟直接在vLLM[18]的融合moe內核內部進行。這樣，反量化過程幾乎不增加額外開銷，甚至導致比BF16更低的延遲。論文已經(jīng)將修改后的融合moe內核貢獻給了vLLM。
論文的ExpertsInt8方法具有多個優(yōu)點
?首先，它速度快；量化僅在模型加載時花費幾秒鐘。
?其次，與vLLM中的大多數(shù)其他技術不同，它不依賴于校準，校準可能需要數(shù)小時或數(shù)天，并且可能不穩(wěn)定。
?第三，論文仍然可以使用BF16來保存大型激活。
?第四，它可以在A100 GPU上使用，而FP8僅在H100上可用。
?最后，論文的量化在延遲上與FP8匹配，同時超越其他量化技術，且沒有質量損失。
圖2比較了使用Jamba-1.5-Mini、Jamba-1.5-Large以及兩個Mixtral模型（8x78B和8x22B）的不同量化技術的延遲。在IH100 GPU上，ExpertsInt8與FP8的延遲相匹配。在A100上，由于FP8不可用，ExpertsInt8是一種有吸引力的技術，大大優(yōu)于GPTQ。結合上述ExpertsInt8的優(yōu)勢，這使得它成為服務大型MoE模型的有吸引力的量化技術。
圖2：不同量化技術的比較，展示了在1024個token上下文和128個token解碼條件下的端到端延遲。ExpertsInt8與FP8表現(xiàn)相似，同時快速且易于應用，仍允許BF16激活，并且適用于A100 GPU，而FP8在這些GPU上不可用。
2.3.2 激活損失
在預訓練過程中，論文發(fā)現(xiàn)某些激活，即特定專家的輸出以及最后Mamba層的輸出，對于某些輸入token，其幅度逐漸增加，最終達到高達4 \times 10^9的值。盡管論文沒有發(fā)現(xiàn)這對預訓練本身造成傷害，預訓練是在BF16精度下進行的，但激活的幅度可能在推理過程中引起數(shù)值問題，因為某些量化庫僅支持FP16精度的激活，其最大范圍為64K。
為了緩解這些擔憂，論文增加了一個“激活損失”項，與前向傳播中激活的均方成正比，并有一個可配置的α因子，懲罰較大的激活值。通過實驗論文發(fā)現(xiàn)，即使\alpha值高達至少10^-3，這種輔助損失對訓練也沒有影響。對于Jamba-1.5-Large，論文使用了α=10^-5，這足以將激活值降低到可接受的范圍（最大2K-3K）。此外，添加這種輔助損失幾乎立即減少了激活值，使其僅在訓練結束時添加，而不影響訓練速度和質量。
為了驗證這種方法，論文使用FP16激活對模型進行了全面評估，并獲得了與BF16評估相同的結果，沒有任何NaN/溢出。
2.4吞吐量和延遲分析
得益于混合Jamba架構，論文的Jamba-1.5模型提供了出色的吞吐量和延遲。圖3和4分別展示了Jamba-1.5-Mini和Jamba-1.5-Large的情況。如圖所示，論文的模型在延遲和吞吐量方面遠優(yōu)于類似規(guī)模的模型。它們在長上下文中的優(yōu)勢尤為明顯，存在顯著差距。重要的是，Jamba-1.5-Large即使在長上下文中也能高效運行，而大型LLaMA3-405B在相同硬件上無法運行。
圖3：Jamba-1.5-Mini與其他模型在延遲和吞吐量方面的比較。所有測量均在2個A100 80GB GPU上進行，批量大小為1，輸出長度為512個token。Jamba-1.5-Mini表現(xiàn)出更好的延遲，尤其是在大型上下文中，輸出token吞吐量僅略有下降。
圖4：Jamba-1.5-Large與其他模型在延遲和吞吐量方面的比較。所有測量均在8塊A100 80GB GPU上進行，批量大小為1，輸出長度為512個token。Jamba-1.5-Large在大型上下文中表現(xiàn)出更好的延遲，輸出token吞吐量僅略有下降。LLaMA-3.1-405B的結果截斷至64K，因為該模型在8塊80GB GPU上無法適應超過100Ktoken的上下文長度。
2.5 訓練
2.5.1 訓練基礎設施和數(shù)據(jù)
Jamba-1.5-Large在NVIDIA H100 GPU上使用論文自有的專有框架進行訓練，該框架包括FSDP、張量并行、序列并行和專家并行。對于后者，論文采用了MegaBlocks
2.5.2訓練階段
該模型分三個階段進行訓練。在預訓練階段，首先在2024年3月更新的自有數(shù)據(jù)集上進行訓練。論文的預訓練數(shù)據(jù)集是公開可用的
網(wǎng)頁文檔、代碼、書籍和科學文章的混合體。
論文的預處理流程包括解析、質量過濾和去重。為了充分利用公開可用數(shù)據(jù)，論文開發(fā)了自己的自有解析器，并使用它提取文本和格式。確切的數(shù)據(jù)混合是通過各種消融實驗確定的。這一階段包括多語言數(shù)據(jù)，重點是以下語言：
英語、西班牙語、法語、葡萄牙語、意大利語、荷蘭語、德語、阿拉伯語和希伯來語。然后，在中間訓練階段進行了一小段時間的高比例長文檔訓練，以強調其遠程能力。最后，模型進行了Post-training。
2.5.3 Post-training
論文的Post-training方法旨在同時實現(xiàn)兩個目標：(i) 賦予模型各種技能和對話能力；(ii) 保留預訓練尤其是中間訓練的長上下文能力。這兩個目標部分存在沖突，因為大多數(shù)可用的Post-training數(shù)據(jù)集由相對較短的示例組成。
鑒于這些考慮，論文的Post-training過程包括在高質量對話數(shù)據(jù)、技能特定數(shù)據(jù)和長上下文數(shù)據(jù)上的監(jiān)督微調。混合這些不同類型的數(shù)據(jù)旨在保留長上下文能力并獲取所需技能。如以下評估所示，論文發(fā)現(xiàn)論文的模型在長上下文評估中表現(xiàn)非常出色。
在進行監(jiān)督微調時，論文大量使用合成數(shù)據(jù)，這在最近的基石模型中很常見，并且反映了論文構建結構化數(shù)據(jù)以構建復合AI系統(tǒng)的方法。論文開發(fā)了多個不同的數(shù)據(jù)合成流程，針對不同的模型能力。所有流程都采用以下模式：（i）在目標分布中采樣或生成提示；（ii）從語言模型生成響應；（iii）根據(jù)自動驗證和評分對響應進行質量過濾或排序；以及（iv）后期編輯以去除偽影并適應所需的格式。論文為構成最終數(shù)據(jù)混合的不同數(shù)據(jù)管道使用不同的模型、提示、采樣、過濾和編輯方法。
論文根據(jù)一系列主要是內部的自動指標選擇了最終的訓練配方（數(shù)據(jù)混合和超參數(shù)）。Jamba-1.5模型都使用相同的控制標記和格式模板進行微調，論文將其作為HlF兼容的標記器和聊天模板的一部分提供；詳見模型卡。
論文提供了幾個合成數(shù)據(jù)生成的顯著例子：
基于表格的問答。論文生成表格數(shù)據(jù)和伴隨的問答對，如論文在表格理解工作[20]中所展示的。然后，論文使用語言模型將表格轉換為自然語言段落。論文生成的訓練示例包括針對給定表格中特定行或列的文本的提取、聚合和歸因任務。
文檔問答。給定一個文檔，論文引導語言模型生成單段落和多段落的問題-答案對。有時，論文通過添加類似文本來嵌入這些示例于更長的上下文中，以鼓勵帶有歸屬的長上下文理解。
工具使用。論文以開源的Glaive函數(shù)調用數(shù)據(jù)集為起點，通過各種啟發(fā)式方法和輸出模式的驗證進行過濾。為了支持并行函數(shù)調用，論文首先為Glaive中的每個函數(shù)生成多個有效的參數(shù)分配。接著，論文從相同函數(shù)和不同函數(shù)中抽取這些有效參數(shù)分配的子集，以生成對應函數(shù)調用集合的用戶請求。最后，論文引導一個函數(shù)調用語言模型響應這些生成的用戶請求，并僅保留函數(shù)調用匹配原始參數(shù)分配的響應。
可引導性。論文定義了一組易于驗證的指令，并合成了包含通用文檔草擬任務及一個或多個約束條件的提示。論文從語言模型中生成這些提示的完成結果，并基于細粒度指令的驗證和通用獎勵模型進行拒絕采樣。為了支持系統(tǒng)消息中的指令，論文選擇了多個共享細粒度指令的此類提示。
2.5.4 一些觀察
論文分享了從Jamba-l.5開發(fā)過程中得出的一些觀察。盡管這些觀察并未完全深入探討，但論文希望它們能啟發(fā)社區(qū)進一步研究這些問題。
首先，盡管論文僅包含了一小部分非英語數(shù)據(jù)，且僅針對特定技能在微調階段進行了處理，但論文的Jamba-1.5模型在多種語言上表現(xiàn)相當出色。如前所述，論文在預訓練階段確實包含了多語言數(shù)據(jù)。因此，論文推測模型能夠在主要使用英語進行微調時利用預訓練階段學到的知識。
其次，論文高效的Jamba架構降低了在長上下文上進行微調的成本，使得在給定預算下能夠進行更多實驗。因此，論文能夠在微調階段嘗試多種不同的訓練方案。
最后，盡管像PPO [33]或DPO [29]這樣的偏好調優(yōu)算法改進了模型輸出與人類意圖之間的一致性，但論文發(fā)現(xiàn)，精心生成的合成數(shù)據(jù)、數(shù)據(jù)過濾和監(jiān)督微調的組合對于獲得強大的微調模型至關重要。
2.6 評估
雖然論文認為基準測試僅部分相關于實際應用的成功和用戶滿意度，但論文仍報告了關鍵公共基準的結果。首先，論文報告了標準學術基準的結果。然后，論文在聊天機器人基準上評估模型。最后，論文對Jamba-1.5-Large進行了多項長上下文評估和多語言評估。
論文與近期同尺寸范圍內的開放權重模型進行了比較：與Jamba-1.5-Large相比，有LLaMA-3.1 70B和Mistral-Large-2-123B；與Jamba-1.5-Mini相比，有LLaMA-3.1-8B和Gemma-2-9B。
2.6.1 學術基準
論文報告了一系列標準學術基準的結果：MMLU、MMLU-Pro、GPQA、ARC-Challence、BBII和HumanEval 。論文還評估了IFEval指令遵循數(shù)據(jù)集和BFCL v1函數(shù)調用數(shù)據(jù)集。最后，論文在RealToxicity和TruthfulQA上報告了安全評估結果。
表2將Jamba-1.5-Large與幾個公開可用且規(guī)模相當?shù)哪Ｐ瓦M行了比較。所有結果均來自官方來源或由論文評估，如表中所示。論文觀察到，Jamba-1.5模型在包括知識、推理、指令遵循和功能調用能力在內的標準學術基準上，與近期公開可用的最先進模型表現(xiàn)相當。論文還觀察到與文獻中報告的安全指標相似。
重要的是，如上所述，Jamba-1.5模型在實現(xiàn)這些結果的同時，提供了更好的吞吐量和延遲。
表2：Jamba-1.5模型在獲得與同等規(guī)模模型相似性能的同時，享受到了更好的吞吐量和延遲。
2.6.2 聊天機器人評估
論文評估了Jamba-1.5模型在兩個聊天機器人場景中的表現(xiàn)：Arena-Hard ，一組500個具有挑戰(zhàn)性的用戶查詢，使用GPT4-Turbo作為評判標準，以及WildBench，使用GPT4-Turbo作為評判標準并進行了長度偏差緩解。如表3所示，Jamba-1.5模型在這些評估中取得了優(yōu)異的結果，其中Jamba-1.5-Large超過了LLaMA-3.1 70B，但略遜于Mistral-Large-2 123B，后者擁有大約30%更多的活躍參數(shù)。
表3：Jamba-1.5模型與類似大小模型在聊天機器人基準測試中的比較。Jamba-1.5模型在性能相似的情況下，具有更好的吞吐量和延遲。" 由論文進行的評估。
2.6.3 長上下文評估
發(fā)布的模型能夠處理長達256K個標記的上下文長度。在本節(jié)中，論文對其在測試其長上下文能力的合成和自然主義基準上進行了評估。
2.6.3.1 RULER
論文在RULER基準上進行了評估，這是一組13個合成任務，旨在評估語言模型的長上下文能力。RULER包括8種針在草堆中的檢索任務變體，包括多個‘needles’[2]。它還包括一個變量跟蹤任務，其中應返回一系列變量綁定，兩個聚合任務，其中一個需要返回最常見的單詞，以及兩個問答任務，其中從自然主義數(shù)據(jù)集[30, 41]中插入的段落來模擬長上下文。
結果展示在表4中。在所有公開和專有模型中，Jamba-1.5-Mini和Jamba-1.5-Large是唯一確認有效長度為256Ktoken的模型。Gemini-pro在原始RULER論文中報告了高達128K的良好結果。然而，盡管論文付出了很大努力，仍無法重現(xiàn)這些結果。論文檢查了Gemini-pro的生成內容，發(fā)現(xiàn)該模型經(jīng)常無法回答或生成拒絕。由于官方RULER結果來自預覽版本，論文假設Gemini-pro自那時起經(jīng)歷了更新，這些更新?lián)p害了其在RULER上的性能。
表4：Jamba-1.5模型與其他公開和專有模型在RULER基準上的比較。其他模型的結果來自RULER Github。bigstar 由論文進行的評估。Jamba-1.5模型是唯一確認有效長度為256Ktoken的模型。
2.6.3.2 Infinite-BENCH
接下來，論文在cOBENCH數(shù)據(jù)集上進行評估，該數(shù)據(jù)集旨在評估語言模型的長上下文能力，平均長度為100K個標記。論文重點關注理解長篇小說的兩個英語任務：問答（EN.QA）和多項選擇問答（EN.MC）。如表5所示，Jamba-1.5模型在這種情況下表現(xiàn)非常出色，優(yōu)于同樣大小的LLaMA-3.1和Mistral-Large-2模型。（由于Gemma-2 9B的上下文窗口較短，僅為8K，因此未報告其結果。）
表5：Jamba-1.5模型在長上下文評估中優(yōu)于同樣大小的LLaMA-3和Mistral-Large-2模型。T評估由論文進行。
2.6.4 多語言能力
論文對Jamba-1.5在非英語語言中的能力進行了基本評估。特別是，論文報告了通過LM Evaluation Harness 分發(fā)的多語言MMLU數(shù)據(jù)集上的結果。表6顯示了結果，其中Jamba-1.5-Mini與其比較點相比表現(xiàn)相似或更好。Jamba-1.5-Large略落后于其可比模型，但仍展現(xiàn)出良好的多語言能力。
表6：Jamba-1.5與其他模型在多語言MMLU數(shù)據(jù)集上的比較。
2.7 對齊與安全考量
論文模型對齊的方法是通過在模型行為與客戶期望之間建立透明度來驅動的。論文的模型默認遵循基于論文參與行業(yè)標準機構、智庫以及與客戶直接經(jīng)驗的商業(yè)行為準則。論文認為這是一種持續(xù)且不斷發(fā)展的合作關系。此外，企業(yè)有多種方式來控制模型行為，以反映其獨特的價值觀和文化，例如額外的培訓和微調、系統(tǒng)消息和提示工程?？傮w而言，論文的AI行為準則基于以下目標：
?使模型行為和輸出與公司價值觀和規(guī)范的商業(yè)禮儀相一致。
?明確聲明預期行為的條款，以便錯誤/漏洞易于識別。
?與客戶合作，并將行為映射到他們的最佳實踐。
?持續(xù)收集反饋，以監(jiān)控并積極改進行為。
根據(jù)論文在OECD
任務組中的角色，該任務組旨在為應用G7廣島人工智能系統(tǒng)開發(fā)組織行為準則開發(fā)一個監(jiān)控機制，論文將模型對齊工作與OECD基于價值觀的AI原則相結合：包容性增長、可持續(xù)發(fā)展與福祉；以人為本的價值和公平性；透明度和可解釋性；魯棒性、安全性和安全性；以及問責制。
對于前四個原則，論文詳細闡述了行為預期或準則，并提供了可用于訓練/對齊和測試合規(guī)性的示例。問責原則側重于Al21在承擔模型行為責任中的角色。論文認為，這種問責主要通過與客戶、監(jiān)管機構和獨立第三方的透明度和接觸來體現(xiàn)。論文與經(jīng)合組織（OECD）、斯坦福大學的HELM [23]和FMTI [3]以及此類文件的接觸，展示了這一承諾，以及論文在FMTI中的高排名（截至2024年5月排名第二）。
論文創(chuàng)建了60個準則，這些準則與OECD原則相對應。這些準則被表述為論文的模型應避免的行為指令。完整列表將公開發(fā)布。
