欢迎 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 架构并将其开源。
