新智元報道
今天，阿布扎比支持的技術創新研究所(TII) 開源了全球第一個通用的大型Mamba架構模型——Falcon Mamba 7B。
雖然之前Mistral已經發過Mamba架構的Codestral Mamba模型，但僅針對編碼；Falcon Mamba則是通用模型，能夠處理各種文本生成任務。
它是繼Falcon 180B、Falcon 40B和Falcon 2之后TII的第四個開放模型，與Falcon系列之前的型號不同，Falcon Mamba 7B完全采用SSLM架構而不是傳統的Transformer架構。
Mamba架構橫空出世后，體現出了內存效率方面的顯著優勢，無需額外的內存需求即可生成大量文本。
如今，SSLM正在逐漸蠶食Transformer架構原本「大一統」的地位。
測評數據顯示，Falcon Mamba 7B性能已經超越同尺寸級別的領先模型，例如Meta最新開源模型Llama 3.1 8B和Mistral 7B。
Falcon Mamba 7B將根據TII Falcon License 2.0發布，這是一個基于Apache 2.0的許可證，其中包括促進負責任地使用人工智能的使用政策。
Falcon Mamba 7B有什么特別之處？
雖然Transformer模型仍然主導著AI大模型領域，但研究人員指出，該架構在處理較長文本時可能會遇到困難。
Transformer的自注意力機制（Self-Attention）讓模型可以關注輸入序列中的所有位置，并為每個位置分配不同的注意力權重。
這使得模型能夠更好地處理長距離的依賴關系，也就是說，對于句子中距離較遠的單詞，模型也能有效地捕獲其關系。
這種通過比較文本中每個token來理解上下文的方式，需要更多的計算能力和內存來處理不斷增長的上下文窗口。
如果資源沒有相應擴展，推理速度會變慢，最終無法處理超過某個固定長度的文本。
為了解決這些難題，狀態空間語言模型（SSLM）架構應運而生，該架構通過在處理單詞時持續更新「狀態」，已成為一種有前途的替代方案。它已經被一些組織部署，TII是最新的采用者。
這個全新的Falcon模型正是使用了CMU和普林斯頓大學的研究人員在2023年12月的一篇論文中最初提出的Mamba SSM架構。
該架構使用選擇機制，使模型能夠根據輸入動態調整其參數。
通過這種方式，模型可以專注于或忽略特定輸入，類似于Transformer中的注意力機制，但同時具備處理長文本序列（如整本書）的能力，而無需額外的內存或計算資源。
TII指出，這種方法使模型適用于企業級機器翻譯、文本摘要、計算機視覺、音頻處理以及估計和預測等任務。
上面提到，基于注意力機制的Transformer是當今所有最強大語言模型中占主導地位的架構。然而，由于計算和內存成本隨著序列長度的增加而增加，注意力機制在處理長序列時存在根本限制。
各種替代架構，特別是SSLM，試圖解決序列擴展限制，但性能不及最先進的Transformer。
Falcon Mamba模型在不損失性能的前提下，可以突破序列擴展限制。
Falcon Mamba基于去年12月提出的第一版Mamba架構，增加了RMS歸一化層以確保在大規模訓練中保持穩定性。
這種架構選擇確保了Falcon Mamba模型：
- 可以在不增加任何內存存儲的情況下處理任意長度的序列，特別是可以在單張A10 24GB GPU上運行；
- 無論上下文大小，生成新token所需的時間恒定。
Falcon Mamba使用約5500GT（相當于5.5B token）的數據進行訓練，主要由RefinedWeb數據組成，并添加了公共來源的高質量技術數據和代碼數據。
在大部分訓練中使用了恒定的學習率，隨后進行了一個較短的學習率衰減階段。
在最后階段，還加入了一小部分高質量的精選數據，以進一步提升模型性能。
性能評估
使用lm-evaluation-harness包對新排行榜版本的所有基準測試進行模型評估，然后使用HuggingFace分數歸一化處理評估結果。
如下圖所示，Falcon Mamba 7B獲得15.04的均分，超過Llama 3.1 8B 13.41分和Mistral 7B 14.50分。
此外，還使用了lighteval對大語言模型排行榜第一版的基準測試進行評估。
可以看到，Falcon Mamba 7B僅次于Transformer架構的Falcon 2 11B，分數仍然超過Gemma、Llama、Mistral等同等規模的知名模型。
理論上來說，SSM模型在處理大規模序列時具有效率優勢。
為了驗證模型的大規模序列處理能力，使用optimum-benchmark庫，對Falcon Mamba和流行的Transformer模型在內存使用和生成吞吐量方面進行了比較。
為了公平比較，將所有Transformer模型的詞匯大小調整為與Falcon Mamba一致，因為這對模型的內存需求有很大影響。
在查看結果之前，先討論序列中提示詞（預填充）和生成（解碼）部分的區別。
預填充的細節對于SSM而言，比對于Transformer模型更為重要。
當Transformer生成下一個token時，它需要關注上下文中所有先前token的鍵和值，這意味著內存需求和生成時間都會隨著上下文長度線性增長。
而SSM僅關注并存儲其遞歸狀態，因此在生成大規模序列時不需要額外的內存或時間。
雖然這解釋了SSM在解碼階段相對于Transformer的優勢，但在預填充階段需要使用新方法來充分利用SSM架構。
預填充的標準方法是并行處理整個提示詞以充分利用GPU。這種方法在optimum-benchmark庫中使用，我們稱之為并行預填充。
并行預填充需要將提示詞每個token的隱藏狀態存儲在內存中。對于Transformer，這額外的內存主要由存儲的KV緩存占據。
對于SSM模型，不需要緩存，存儲隱藏狀態的內存成為唯一與提示詞長度成比例的部分。
因此，內存需求將隨提示詞長度增長，SSM模型將失去處理任意長序列的能力，類似于Transformer。
并行預填充的替代方法是逐個處理token提示詞，我們稱之為順序預填充。
類似于序列并行處理，它也可以大規模地處理提示詞，而不是單個token，以更好地利用GPU。
雖然順序預填充對Transformer意義不大，但它為SSM模型帶來了處理任意長提示詞的可能性。
考慮到這些觀點，實驗首先測試了可以在單個24GB A10 GPU上適應的最大序列長度。
其中，批大小固定為1，使用float32精度。
即使在并行預填充中，Falcon Mamba也能適應比Transformer更大的序列，而在順序預填充中發揮了全部潛力，可以處理任意長度的提示詞。
接下來，在提示詞長度為1，生成token數量最多為130k的情況下測量生成吞吐量，使用批大小為1，并在H100 GPU上進行。
結果如圖所示。可以觀察到，Falcon Mamba在生成所有token時保持恒定的吞吐量，且GPU峰值內存沒有增加。
而對于Transformer模型，隨著生成token數量的增加，峰值內存增加，生成速度變慢。
如何使用？
Falcon Mamba架構將在HuggingFace transformers庫的下一個版本（4.45.0以上）中提供。
使用Falcon Mamba 7B模型，需要安裝最新版本的HuggingFace transformers，或從源代碼安裝庫。
Falcon Mamba與HuggingFace提供的大多數API兼容，這些API已經比較熟悉，例如：
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內存較小的情況下運行模型，例如：
此外，還推出了Falcon Mamba的指令微調版本，該版本經過額外50億個token的監督微調（SFT），這種擴展訓練提高了模型在執行指令任務時的精確性和有效性。
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_config=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))可以通過演示體驗該指令模型的功能，對于聊天模板，可以使用以下格式：
<|im_start|>user
prompt<|im_end|>
<|im_start|>assistant用戶還可以直接使用基礎模型和指令模型的4-bit轉換版本，但要保證GPU與bitsandbytes庫兼容，才能運行量化模型。
用戶還可以通過torch.compile獲得更快的推理速度，加載模型后，只需調用model = torch.compile(model)。
