Falcon Mamba 7B Release Notes

Falcon Mamba 7B is a new open-access model from the Technology Innovation Institute (TII) that demonstrates that State Space Language Models (SSLMs) can achieve performance competitive with state-of-the-art (SoTA) transformers without the linear increase in compute and memory costs associated with the attention mechanism.

Pure Mamba Architecture and Sequence Scaling

Falcon Mamba is based on the original Mamba architecture, utilizing selective state spaces to overcome the sequence scaling limitations of Transformers. To ensure stable training at scale, TII added extra RMS normalization layers to the base Mamba design.

This architecture provides two primary technical advantages over attention-based models:

  • Arbitrary Sequence Length: The model can process sequences of any length without a corresponding increase in memory storage, allowing it to fit on a single A10 24GB GPU.
  • Constant Generation Time: The time required to generate a new token remains constant regardless of the size of the context.

Model Training and Data

Falcon Mamba was trained on approximately 5,500GT of data. The dataset primarily consists of RefinedWeb data, supplemented by high-quality technical data and code from public sources. The training process employed a constant learning rate for the majority of the duration, followed by a short learning rate decay stage that incorporated a small portion of high-quality curated data to boost final performance.

Performance Benchmarks

Evaluations conducted via lm-evaluation-harness (with Hugging Face score normalization) and lighteval show that Falcon Mamba-7B is the strongest pure SSM model to date, often competing with or exceeding the performance of hybrid SSM-attention models and several transformer-based models.

New Leaderboard Version Results

Model Name IFEval BBH MATH LvL5 GPQA MUSR MMLU-PRO Average
Falcon Mamba-7B 33.36 19.88 3.63 8.05 10.86 14.47 15.04
recurrentgemma-9b 30.76 14.80 4.83 4.70 6.60 17.88 13.20
Falcon2-11B 32.61 21.94 2.34 2.80 7.53 15.44 13.78
Meta-Llama-3.1-8B 12.70 25.29 4.61 6.15 8.98 24.95 13.78
gemma-7B 26.59 21.12 6.42 4.92 10.98 21.64 15.28

LLM Leaderboard (First Version) Results

Model Name ARC HellaSwag MMLU Winogrande TruthfulQA GSM8K Average
Falcon Mamba-7B 62.03 80.82 62.11 73.64 53.42 52.54 64.09
recurrentgemma-9b 52.00 80.40 60.50 73.60 38.60 42.60 57.95
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

Memory and Throughput Analysis

Falcon Mamba exhibits significant efficiency gains over transformer models during both the prefill and decode stages of sequence processing.

Prefill Efficiency

There are two primary methods for processing prompts (prefill):

  1. Parallel Prefill: Processes the entire prompt in parallel to maximize GPU utilization. In this mode, memory requirements scale with prompt length due to the need to store hidden states. Even so, Falcon Mamba fits larger sequences than transformers.
  2. Sequential Prefill: Processes the prompt token-by-token or in chunks. This method allows SSMs to process arbitrary long prompts without the memory scaling issues found in transformers.

Generation Throughput

In tests using an H100 GPU with a prompt length of 1 and up to 130k generated tokens, Falcon Mamba maintained constant throughput and stable CUDA peak memory. In contrast, transformer models showed increasing peak memory usage and decreasing generation speed as the sequence length grew.

Implementation and Availability

Falcon Mamba is integrated into the Hugging Face transformers library (version >4.45.0) and is compatible with AutoModelForCausalLM and pipeline APIs.

Model Variants

  • Base Model: The standard pre-trained model.
  • Instruct Model: A version fine-tuned with 5 billion tokens of supervised fine-tuning (SFT) data for improved instructional task performance.
  • Quantized Versions: 4-bit converted versions of both base and instruct models are available for users with bitsandbytes compatible GPUs.

Optimization

Users can achieve faster inference by applying torch.compile(model) after loading the model.

Sources