Falcon Mamba 是由阿布扎比的 Technology Innovation Institute (TII) 開發並基於 TII Falcon Mamba 7B License 1.0 的開放獲取模型。該模型是開放獲取的，所以任何人都可以在 Hugging Face 生態系統中 這裡 使用它進行研究或應用。
在這篇部落格中，我們將深入模型的設計決策、探究模型與其他現有的 SoTA 模型相比的競爭力，以及如何在 Hugging Face 生態系統中使用它。
第一款通用的大規模純 Mamba 模型
目前，所有頂級大型語言模型都使用基於注意力機制的 Transformer 架構。然而，由於計算和記憶體成本隨序列長度增加而增加，注意力機制在處理大序列時存在根本性的侷限性。狀態空間語言模型 (SSLMs) 等各種替代架構嘗試解決序列擴充套件的限制，但在效能上仍不及最先進的 Transformer 架構模型。
透過 Falcon Mamba，我們證明了序列擴充套件的限制確實可以在不損失效能的情況下克服。Falcon Mamba 基於原始的 Mamba 架構，該架構在 Mamba: Linear-Time Sequence Modeling with Selective State Spaces 中提出，並增加了額外的 RMS 標準化層以確保大規模穩定訓練。這種架構選擇確保 Falcon Mamba:
- 能夠處理任意長度的序列，而不增加記憶體儲存，特別是適用於單個 A10 24GB GPU。
- 生成新令牌的時間是恆定的，不論上下文的大小。
模型訓練
Falcon Mamba 訓練所用的資料量約為 5500GT，主要包括經過精選的網路資料，並補充了來自公開源的高質量技術和程式碼資料。我們在大部分訓練過程中使用恆定的學習率，隨後進行了一個相對較短的學習率衰減階段。在最後這個階段，我們還新增了一小部分高質量的策劃資料，以進一步提高模型效能。
模型評估
我們使用 lm-evaluation-harness 包在新排行榜版本的所有基準上評估我們的模型，然後使用 Hugging Face 分數規範化方法規範化評估結果。model name``IFEval``BBH``MATH LvL5``GPQA``MUSR``MMLU-PRO``Average
model name |
IFEval |
BBH |
MATH LvL5 |
GPQA |
MUSR |
MMLU-PRO |
Average |
|---|---|---|---|---|---|---|---|
| Pure SSM models | |||||||
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 |
| Hybrid SSM-attention models | |||||||
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 models | |||||||
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 排行榜第一版的基準測試上對模型進行了評估。model name``ARC``HellaSwag``MMLU``Winogrande``TruthfulQA``GSM8K``Average
model name |
ARC |
HellaSwag |
MMLU |
Winogrande |
TruthfulQA |
GSM8K |
Average |
|---|---|---|---|---|---|---|---|
| Pure SSM models | |||||||
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 |
| Hybrid SSM-attention models | |||||||
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 models | |||||||
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，因為這對模型的記憶體需求有很大影響。
在介紹結果之前，首先討論提示 (prefill) 和生成 (decode) 部分序列的區別。我們將看到，對於狀態空間模型而言，prefill 的細節比 Transformer 模型更為重要。當 Transformer 生成下一個令牌時，它需要關注上下文中所有之前令牌的鍵和值。這意味著記憶體需求和生成時間都隨上下文長度線性增長。狀態空間模型僅關注並儲存其迴圈狀態，因此不需要額外的記憶體或時間來生成大序列。雖然這解釋了 SSM 在解碼階段相對於 Transformer 的優勢，但 prefill 階段需要額外努力以充分利用 SSM 架構。
prefill 的標準方法是並行處理整個提示，以充分利用 GPU。這種方法在 optimum-benchmark 庫中被使用，並被我們稱為並行 prefill。並行 prefill 需要在記憶體中儲存提示中每個令牌的隱藏狀態。對於 Transformer，這額外的記憶體主要由儲存的 KV 快取所佔據。對於 SSM 模型，不需要快取，儲存隱藏狀態的記憶體成為與提示長度成比例的唯一組成部分。結果，記憶體需求將隨提示長度擴充套件，SSM 模型將失去處理任意長序列的能力，與 Transformer 類似。
另一種 prefill 方法是逐令牌處理提示，我們將其稱為 順序 prefill 。類似於序列並行性，它也可以在更大的提示塊上執行，而不是單個令牌，以更好地利用 GPU。雖然對於 Transformer 來說，順序 prefill 幾乎沒有意義，但它讓 SSM 模型重新獲得了處理任意長提示的可能性。
基於這些考慮，我們首先測試了單個 24GB A10 GPU 可以支援的最大序列長度，具體結果請見下方的圖表。批處理大小固定為 1，我們使用 float32 精度。即使對於並行 prefill，Falcon Mamba 也能適應比 Transformer 更大的序列，而在順序 prefill 中，它釋放了全部潛力，可以處理任意長的提示。
接下來，我們在提示長度為 1 且生成高達 130k 令牌的設定中測量生成吞吐量，使用批次大小 1 和 H100 GPU。結果報告在下方的圖表中。我們觀察到，我們的 Falcon Mamba 在恆定的吞吐量下生成所有令牌，且 CUDA 峰值記憶體沒有增加。對於 Transformer 模型，峰值記憶體隨生成令牌數的增加而增長，生成速度也隨之減慢。
接下來，我們在使用單個 H100 GPU 和批次大小為 1 的設定中，測量了提示長度為 1 且生成高達 130,000 個令牌的生成吞吐量。結果顯示在下方的圖表中。我們觀察到，我們的 Falcon Mamba 能夠以恆定的吞吐量生成所有令牌，並且 CUDA 峰值記憶體沒有任何增加。對於 Transformer 模型，隨著生成令牌數量的增加，峰值記憶體增長，生成速度減慢。
在 Hugging Face transformers 中如何使用 Falcon Mamba？
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 位轉換版本。確保您有權訪問與 bitsandbytes 庫相容的 GPU 來執行量化模型。
您還可以使用 torch.compile 實現更快的推理; 只需在載入模型後呼叫 model = torch.compile(model) 。
致謝
我們感謝 Hugging Face 團隊在整合過程中提供的無縫支援，特別鳴謝以下人員:
- Alina Lozovskaya 和 Clementine Fourrier 幫助我們在排行榜上評估模型
- Arthur Zucker 負責 transformers 的整合
- Vaibhav Srivastav, hysts 和 Omar Sanseviero 在 Hub 相關問題上提供的支援
作者還要感謝 Tri Dao 和 Albert Gu 將 Mamba 架構實現並開源給社群。
英文原文: https://hf.co/blog/falconmamba
原文作者: Jingwei Zuo, Maksim Velikanov, Rhaiem, Ilyas Chahed, Younes Belkada, Guillaume Kunsch
譯者: Evinci