Mamba再次挑戰(zhàn)霸主Transformer！首個通用Mamba開源大模型一鳴驚人
今天，阿布扎比支持的技術(shù)創(chuàng)新研究所(TII) 開源了全球第一個通用的大型Mamba架構(gòu)模型——Falcon Mamba 7B。
雖然之前Mistral已經(jīng)發(fā)過Mamba架構(gòu)的Codestral Mamba模型，但僅針對編碼；Falcon Mamba則是通用模型，能夠處理各種文本生成任務(wù)。
它是繼Falcon 180B、Falcon 40B和Falcon 2之后TII的第四個開放模型，與Falcon系列之前的型號不同，F(xiàn)alcon Mamba 7B完全采用SSLM架構(gòu)而不是傳統(tǒng)的Transformer架構(gòu)。
圖片
Mamba架構(gòu)橫空出世后，體現(xiàn)出了內(nèi)存效率方面的顯著優(yōu)勢，無需額外的內(nèi)存需求即可生成大量文本。
如今，SSLM正在逐漸蠶食Transformer架構(gòu)原本「大一統(tǒng)」的地位。
測評數(shù)據(jù)顯示，F(xiàn)alcon Mamba 7B性能已經(jīng)超越同尺寸級別的領(lǐng)先模型，例如Meta最新開源模型Llama 3.1 8B和Mistral 7B。
Falcon Mamba 7B將根據(jù)TII Falcon License 2.0發(fā)布，這是一個基于Apache 2.0的許可證，其中包括促進(jìn)負(fù)責(zé)任地使用人工智能的使用政策。
Falcon Mamba 7B有什么特別之處？
雖然Transformer模型仍然主導(dǎo)著AI大模型領(lǐng)域，但研究人員指出，該架構(gòu)在處理較長文本時可能會遇到困難。
Transformer的自注意力機(jī)制（Self-Attention）讓模型可以關(guān)注輸入序列中的所有位置，并為每個位置分配不同的注意力權(quán)重。
這使得模型能夠更好地處理長距離的依賴關(guān)系，也就是說，對于句子中距離較遠(yuǎn)的單詞，模型也能有效地捕獲其關(guān)系。
這種通過比較文本中每個token來理解上下文的方式，需要更多的計算能力和內(nèi)存來處理不斷增長的上下文窗口。
如果資源沒有相應(yīng)擴(kuò)展，推理速度會變慢，最終無法處理超過某個固定長度的文本。
為了解決這些難題，狀態(tài)空間語言模型（SSLM）架構(gòu)應(yīng)運(yùn)而生，該架構(gòu)通過在處理單詞時持續(xù)更新「狀態(tài)」，已成為一種有前途的替代方案。它已經(jīng)被一些組織部署，TII是最新的采用者。
這個全新的Falcon模型正是使用了CMU和普林斯頓大學(xué)的研究人員在2023年12月的一篇論文中最初提出的Mamba SSM架構(gòu)。
圖片
論文地址：https://arxiv.org/pdf/2312.00752
該架構(gòu)使用選擇機(jī)制，使模型能夠根據(jù)輸入動態(tài)調(diào)整其參數(shù)。
通過這種方式，模型可以專注于或忽略特定輸入，類似于Transformer中的注意力機(jī)制，但同時具備處理長文本序列（如整本書）的能力，而無需額外的內(nèi)存或計算資源。
TII指出，這種方法使模型適用于企業(yè)級機(jī)器翻譯、文本摘要、計算機(jī)視覺、音頻處理以及估計和預(yù)測等任務(wù)。
首個通用大規(guī)模Mamba模型
上面提到，基于注意力機(jī)制的Transformer是當(dāng)今所有最強(qiáng)大語言模型中占主導(dǎo)地位的架構(gòu)。然而，由于計算和內(nèi)存成本隨著序列長度的增加而增加，注意力機(jī)制在處理長序列時存在根本限制。
各種替代架構(gòu)，特別是SSLM，試圖解決序列擴(kuò)展限制，但性能不及最先進(jìn)的Transformer。
Falcon Mamba模型在不損失性能的前提下，可以突破序列擴(kuò)展限制。
Falcon Mamba基于去年12月提出的第一版Mamba架構(gòu)，增加了RMS歸一化層以確保在大規(guī)模訓(xùn)練中保持穩(wěn)定性。
這種架構(gòu)選擇確保了Falcon Mamba模型：
- 可以在不增加任何內(nèi)存存儲的情況下處理任意長度的序列，特別是可以在單張A10 24GB GPU上運(yùn)行；
- 無論上下文大小，生成新token所需的時間恒定。
模型訓(xùn)練
Falcon Mamba使用約5500GT（相當(dāng)于5.5B token）的數(shù)據(jù)進(jìn)行訓(xùn)練，主要由RefinedWeb數(shù)據(jù)組成，并添加了公共來源的高質(zhì)量技術(shù)數(shù)據(jù)和代碼數(shù)據(jù)。
在大部分訓(xùn)練中使用了恒定的學(xué)習(xí)率，隨后進(jìn)行了一個較短的學(xué)習(xí)率衰減階段。
在最后階段，還加入了一小部分高質(zhì)量的精選數(shù)據(jù)，以進(jìn)一步提升模型性能。
性能評估
使用lm-evaluation-harness包對新排行榜版本的所有基準(zhǔn)測試進(jìn)行模型評估，然后使用HuggingFace分?jǐn)?shù)歸一化處理評估結(jié)果。
如下圖所示，F(xiàn)alcon Mamba 7B獲得15.04的均分，超過Llama 3.1 8B 13.41分和Mistral 7B 14.50分。
圖片
此外，還使用了lighteval對大語言模型排行榜第一版的基準(zhǔn)測試進(jìn)行評估。
可以看到，F(xiàn)alcon Mamba 7B僅次于Transformer架構(gòu)的Falcon 2 11B，分?jǐn)?shù)仍然超過Gemma、Llama、Mistral等同等規(guī)模的知名模型。
處理大規(guī)模序列
理論上來說，SSM模型在處理大規(guī)模序列時具有效率優(yōu)勢。
為了驗(yàn)證模型的大規(guī)模序列處理能力，使用optimum-benchmark庫，對Falcon Mamba和流行的Transformer模型在內(nèi)存使用和生成吞吐量方面進(jìn)行了比較。
為了公平比較，將所有Transformer模型的詞匯大小調(diào)整為與Falcon Mamba一致，因?yàn)檫@對模型的內(nèi)存需求有很大影響。
在查看結(jié)果之前，先討論序列中提示詞（預(yù)填充）和生成（解碼）部分的區(qū)別。
預(yù)填充的細(xì)節(jié)對于SSM而言，比對于Transformer模型更為重要。
當(dāng)Transformer生成下一個token時，它需要關(guān)注上下文中所有先前token的鍵和值，這意味著內(nèi)存需求和生成時間都會隨著上下文長度線性增長。
而SSM僅關(guān)注并存儲其遞歸狀態(tài)，因此在生成大規(guī)模序列時不需要額外的內(nèi)存或時間。
圖片
雖然這解釋了SSM在解碼階段相對于Transformer的優(yōu)勢，但在預(yù)填充階段需要使用新方法來充分利用SSM架構(gòu)。
預(yù)填充的標(biāo)準(zhǔn)方法是并行處理整個提示詞以充分利用GPU。這種方法在optimum-benchmark庫中使用，我們稱之為并行預(yù)填充。
并行預(yù)填充需要將提示詞每個token的隱藏狀態(tài)存儲在內(nèi)存中。對于Transformer，這額外的內(nèi)存主要由存儲的KV緩存占據(jù)。
對于SSM模型，不需要緩存，存儲隱藏狀態(tài)的內(nèi)存成為唯一與提示詞長度成比例的部分。
因此，內(nèi)存需求將隨提示詞長度增長，SSM模型將失去處理任意長序列的能力，類似于Transformer。
并行預(yù)填充的替代方法是逐個處理token提示詞，我們稱之為順序預(yù)填充。
類似于序列并行處理，它也可以大規(guī)模地處理提示詞，而不是單個token，以更好地利用GPU。
雖然順序預(yù)填充對Transformer意義不大，但它為SSM模型帶來了處理任意長提示詞的可能性。
考慮到這些觀點(diǎn)，實(shí)驗(yàn)首先測試了可以在單個24GB A10 GPU上適應(yīng)的最大序列長度。
其中，批大小固定為1，使用float32精度。
即使在并行預(yù)填充中，F(xiàn)alcon Mamba也能適應(yīng)比Transformer更大的序列，而在順序預(yù)填充中發(fā)揮了全部潛力，可以處理任意長度的提示詞。
圖片
接下來，在提示詞長度為1，生成token數(shù)量最多為130k的情況下測量生成吞吐量，使用批大小為1，并在H100 GPU上進(jìn)行。
結(jié)果如圖所示?？梢杂^察到，F(xiàn)alcon Mamba在生成所有token時保持恒定的吞吐量，且GPU峰值內(nèi)存沒有增加。
而對于Transformer模型，隨著生成token數(shù)量的增加，峰值內(nèi)存增加，生成速度變慢。
圖片
如何使用？
Falcon Mamba架構(gòu)將在HuggingFace transformers庫的下一個版本（4.45.0以上）中提供。
使用Falcon Mamba 7B模型，需要安裝最新版本的HuggingFace transformers，或從源代碼安裝庫。
Falcon Mamba與HuggingFace提供的大多數(shù)API兼容，這些API已經(jīng)比較熟悉，例如：
from transformers import AutoModelForCausalLM, AutoTokenizer
model_id = "tiiuae/falcon-mamba-7b"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto", device_map="auto")
inputs = tokenizer("Hello world, today", return_tensors="pt").to(0)
output = model.generate(**inputs, max_new_tokens=100, do_sample=True)
print(tokenizer.decode(Output[0], skip_special_tokens=True))它還支持例如bitsandbytes庫量化這樣的功能，以便在GPU內(nèi)存較小的情況下運(yùn)行模型，例如：
此外，還推出了Falcon Mamba的指令微調(diào)版本，該版本經(jīng)過額外50億個token的監(jiān)督微調(diào)（SFT），這種擴(kuò)展訓(xùn)練提高了模型在執(zhí)行指令任務(wù)時的精確性和有效性。
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
model_id = "tiiuae/falcon-mamba-7b"
tokenizer = AutoTokenizer.from_pretrained(model_id)
quantization_config = BitsAndBytesConfig(load_in_4bit=True)
model = AutoModelForCausalLM.from_pretrained(model_id, quantization_cnotallow=quantization_config)
inputs = tokenizer("Hello world, today", return_tensors="pt").to(0)
output = model.generate(**inputs, max_new_tokens=100, do_sample=True)
print(tokenizer.decode(output[0], skip_special_tokens=True))可以通過演示體驗(yàn)該指令模型的功能，對于聊天模板，可以使用以下格式：
<|im_start|>user
prompt<|im_end|>
<|im_start|>assistant用戶還可以直接使用基礎(chǔ)模型和指令模型的4-bit轉(zhuǎn)換版本，但要保證GPU與bitsandbytes庫兼容，才能運(yùn)行量化模型。
用戶還可以通過torch.compile獲得更快的推理速度，加載模型后，只需調(diào)用model = torch.compile(model)。
參考資料：
https://huggingface.co/blog/falconmamba
https://venturebeat.com/ai/falcon-mamba-7bs-powerful-new-ai-architecture-offers-alternative-to-transformer-models/
https://medium.com/@puneetthegde22/mamba-architecture-a-leap-forward-in-sequence-modeling-370dfcbfe44a
