歡迎 FalconMamba：首個強大的無注意力 7B 模型
Falcon Mamba 是由阿布扎比 技術創新研究所 (TII) 釋出的新模型，採用 TII Falcon Mamba 7B 許可 1.0 釋出。該模型開放訪問，可在 Hugging Face 生態系統中的此處獲取，供任何人用於研究或應用目的。
在本部落格中，我們將詳細介紹模型背後的設計決策，模型如何與現有最先進模型競爭，以及如何在 Hugging Face 生態系統中使用它。
首個通用型大規模純 Mamba 模型
基於注意力機制的 Transformer 模型是當今所有最強大大型語言模型的主流架構。然而，由於計算和記憶體成本隨序列長度增加，注意力機制在處理大型序列時存在根本性限制。各種替代架構，特別是狀態空間語言模型 (SSLM)，試圖解決序列縮放限制，但在效能方面卻落後於最先進的 Transformer 模型。
透過 Falcon Mamba，我們證明了在不損失效能的情況下，確實可以克服序列縮放限制。Falcon Mamba 基於原始的 Mamba 架構，該架構在《Mamba：基於選擇性狀態空間的線性時間序列建模》中提出，並額外添加了 RMS 歸一化層，以確保大規模訓練的穩定性。這種架構選擇確保了 Falcon Mamba
- 可以處理任意長度的序列，而無需增加記憶體儲存，特別是可以在單個 A10 24GB GPU 上執行。
- 生成新標記所需的時間是恆定的，無論上下文大小如何（請參閱本節）
模型訓練
Falcon Mamba 使用約 5500 GT 資料進行訓練，主要由 RefinedWeb 資料組成，並增加了來自公共來源的高質量技術資料和程式碼資料。我們大部分訓練使用恆定學習率，隨後是相對較短的學習率衰減階段。在最後階段，我們還添加了一小部分高質量的精選資料，以進一步提高模型效能。
評估
我們使用 lm-evaluation-harness 包在所有新排行榜版本的基準上評估了我們的模型，然後使用 Hugging Face 分數歸一化來歸一化評估結果。
模型名稱 |
IFEval |
BBH |
MATH LvL5 |
GPQA |
MUSR |
MMLU-PRO |
平均分 |
|---|---|---|---|---|---|---|---|
| 純 SSM 模型 | |||||||
Falcon Mamba-7B |
33.36 | 19.88 | 3.63 | 8.05 | 10.86 | 14.47 | 15.04 |
TRI-ML/mamba-7b-rw* |
22.46 | 6.71 | 0.45 | 1.12 | 5.51 | 1.69 | 6.25 |
| 混合 SSM-注意力模型 | |||||||
recurrentgemma-9b |
30.76 | 14.80 | 4.83 | 4.70 | 6.60 | 17.88 | 13.20 |
Zyphra/Zamba-7B-v1* |
24.06 | 21.12 | 3.32 | 3.03 | 7.74 | 16.02 | 12.55 |
| Transformer 模型 | |||||||
Falcon2-11B |
32.61 | 21.94 | 2.34 | 2.80 | 7.53 | 15.44 | 13.78 |
Meta-Llama-3-8B |
14.55 | 24.50 | 3.25 | 7.38 | 6.24 | 24.55 | 13.41 |
Meta-Llama-3.1-8B |
12.70 | 25.29 | 4.61 | 6.15 | 8.98 | 24.95 | 13.78 |
Mistral-7B-v0.1 |
23.86 | 22.02 | 2.49 | 5.59 | 10.68 | 22.36 | 14.50 |
Mistral-Nemo-Base-2407 (12B) |
16.83 | 29.37 | 4.98 | 5.82 | 6.52 | 27.46 | 15.08 |
gemma-7B |
26.59 | 21.12 | 6.42 | 4.92 | 10.98 | 21.64 | 15.28 |
此外，我們還使用 lighteval 評估了 LLM 排行榜第一個版本的基準上的模型。
模型名稱 |
ARC |
HellaSwag |
MMLU |
Winogrande |
TruthfulQA |
GSM8K |
平均分 |
|---|---|---|---|---|---|---|---|
| 純 SSM 模型 | |||||||
Falcon Mamba-7B* |
62.03 | 80.82 | 62.11 | 73.64 | 53.42 | 52.54 | 64.09 |
TRI-ML/mamba-7b-rw* |
51.25 | 80.85 | 33.41 | 71.11 | 32.08 | 4.70 | 45.52 |
| 混合 SSM-注意力模型 | |||||||
recurrentgemma-9b** |
52.00 | 80.40 | 60.50 | 73.60 | 38.60 | 42.60 | 57.95 |
Zyphra/Zamba-7B-v1* |
56.14 | 82.23 | 58.11 | 79.87 | 52.88 | 30.78 | 60.00 |
| Transformer 模型 | |||||||
Falcon2-11B |
59.73 | 82.91 | 58.37 | 78.30 | 52.56 | 53.83 | 64.28 |
Meta-Llama-3-8B |
60.24 | 82.23 | 66.70 | 78.45 | 42.93 | 45.19 | 62.62 |
Meta-Llama-3.1-8B |
58.53 | 82.13 | 66.43 | 74.35 | 44.29 | 47.92 | 62.28 |
Mistral-7B-v0.1 |
59.98 | 83.31 | 64.16 | 78.37 | 42.15 | 37.83 | 60.97 |
gemma-7B |
61.09 | 82.20 | 64.56 | 79.01 | 44.79 | 50.87 | 63.75 |
對於標有*的模型，我們進行了內部任務評估，而對於標有**的模型，結果取自論文或模型卡片。
處理長序列
根據 SSM 模型在處理長序列方面的理論效率，我們使用 optimum-benchmark 庫對 Falcon Mamba 和流行 Transformer 模型之間的記憶體使用和生成吞吐量進行了比較。為了公平比較，我們重新縮放了所有 Transformer 模型的詞彙量，使其與 Falcon Mamba 匹配，因為詞彙量對模型的記憶體需求有很大影響。
在介紹結果之前，我們首先討論序列的提示（預填充）部分和生成（解碼）部分之間的差異。正如我們將看到的，預填充的細節對於狀態空間模型比對於 Transformer 模型更重要。當 Transformer 生成下一個標記時，它需要關注上下文中所有先前標記的鍵和值。這意味著記憶體需求和生成時間都隨上下文長度線性擴充套件。狀態空間模型只關注並存儲其迴圈狀態，因此不需要額外的記憶體或時間來生成長序列。雖然這解釋了 SSM 在解碼階段相對於 Transformer 的優勢，但預填充階段需要額外的努力才能充分利用 SSM 架構。
預填充的標準方法是並行處理整個提示以充分利用 GPU。這種方法在 optimum-benchmark 庫中使用，我們將其稱為並行預填充。並行預填充需要將提示中每個標記的隱藏狀態儲存在記憶體中。對於 Transformer，這種額外的記憶體主要由儲存的 KV 快取的記憶體主導。對於 SSM 模型，不需要快取，儲存隱藏狀態的記憶體成為唯一與提示長度成比例的元件。結果，記憶體需求將隨提示長度縮放，SSM 模型將失去處理任意長序列的能力，類似於 Transformer。
並行預填充的替代方法是逐個標記處理提示，我們將其稱為順序預填充。類似於序列並行性，它也可以在更大的提示塊而不是單個標記上完成，以更好地利用 GPU。雖然順序預填充對於 Transformer 意義不大，但它使 SSM 模型能夠處理任意長的提示。
考慮到這些注意事項，我們首先測試可以容納在單個 24 GB A10 GPU 上的最大序列長度，結果顯示在下面的圖中。批處理大小固定為 1，我們使用 float32 精度。即使對於並行預填充，Falcon Mamba 也可以容納比 Transformer 更長的序列，而在順序預填充中，它充分發揮了潛力，可以處理任意長的提示。
接下來，我們測量了提示長度為 1 且生成標記多達 130k 的設定中的生成吞吐量，使用批處理大小為 1 和 H100 GPU。結果顯示在下面的圖中。我們觀察到我們的 Falcon Mamba 以恆定的吞吐量生成所有標記，並且 CUDA 峰值記憶體沒有任何增加。對於 Transformer 模型，峰值記憶體會增長，並且生成速度會隨著生成標記數量的增加而減慢。
如何在 Hugging Face Transformers 中使用它？
Falcon Mamba 架構將在 Hugging Face Transformers 庫的下一個版本（>4.45.0）中提供。要使用該模型，請確保安裝最新版本的 Hugging Face Transformers 或從原始碼安裝該庫。
Falcon Mamba 相容 Hugging Face 提供的大多數您熟悉的 API，例如 AutoModelForCausalLM 或 pipeline
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 記憶體限制下執行模型，例如
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))
我們也很高興推出 Falcon Mamba 的指令微調版本，該版本已透過額外 50 億個監督微調 (SFT) 資料標記進行微調。這種擴充套件訓練增強了模型以更高的精度和有效性執行指令任務的能力。您可以透過我們的演示體驗指令模型的功能，可在此處獲取。對於聊天模板，我們使用以下格式
<|im_start|>user
prompt<|im_end|>
<|im_start|>assistant
您也可以直接使用基本模型的 4 位轉換版本（在此）和指令模型的 4 位轉換版本（在此）。請確保您擁有與 bitsandbytes 庫相容的 GPU 才能執行量化模型。
您還可以透過使用 torch.compile 來獲得更快的推理速度；一旦載入模型，只需呼叫 model = torch.compile(model) 即可。
致謝
本部落格文章的作者要感謝 Hugging Face 團隊在他們的生態系統中的順利支援和整合，特別是
- Alina Lozovskaya 和 Clementine Fourrier 幫助我們在排行榜上評估模型
- Arthur Zucker 負責 Transformer 整合
- Vaibhav Srivastav、hysts 和 Omar Sanseviero 支援 Hub 相關問題
作者還要感謝 Tri Dao 和 Albert Gu 為社群實現 Mamba 架構並將其開源。
