Falcon Mamba 是阿布扎比技术创新研究所 (TII) 根据 TII Falcon License 2.0 发布的新型号。该模型是开放获取的,可在Hugging Face生态系统中使用,任何人都可以用于他们的研究或应用目的。
在这篇博客中,我们将介绍模型背后的设计决策,该模型相对于其他现有的 SoTA 模型如何具有竞争力,以及如何在 Hugging Face 生态系统中使用它。
第一个通用的大规模纯曼巴模型(pure Mamba model)
基于注意力机制的转换器是当今所有最强大的大型语言模型中使用的主要架构。然而,由于计算和内存成本随着序列长度的增加而增加,注意力机制在处理大型序列时从根本上受到限制。各种替代架构,特别是状态空间语言模型 (SSLM),试图解决序列缩放限制,但与 SoTA 转换器相比,性能有所回落。
通过Falcon Mamba,我们证明了序列缩放限制确实可以克服,而不会降低性能。Falcon Mamba 基于 Mamba 中提出的原始 Mamba 架构:具有选择性状态空间的线性时间序列建模,并添加了额外的 RMS 归一化层以确保大规模的稳定训练。这种架构的选择确保了 Falcon Mamba:
- 可以在不增加内存存储的情况下处理任意长度的序列,特别是适合单个 A10 24GB GPU。
- 无论上下文的大小如何,都需要恒定的时间来生成新的令牌(请参阅此部分)
模型训练
Falcon Mamba 使用 ~ 5500GT 数据进行训练,主要由 RefinedWeb 数据组成,此外还有来自公共来源的高质量技术数据和代码数据。我们在大部分训练中使用恒定学习率,然后是相对较短的学习率衰减阶段。在最后阶段,我们还添加了一小部分高质量的精选数据,以进一步提高模型性能。
评估
我们使用lm-evaluation-harness
软件包在新排行榜版本的所有基准上评估我们的模型,然后使用 Hugging Face 分数归一化对评估结果进行归一化。
model name | IFEval | BBH | MATH LvL5 | GPQA | MUSR | MMLU-PRO | Average |
---|---|---|---|---|---|---|---|
纯 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 |
变压器型号 | |||||||
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 |
---|---|---|---|---|---|---|---|
纯 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 |
变压器型号 | |||||||
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 模型的词汇量以匹配 Falcon Mamba,因为它对模型的内存需求有很大影响。
在讨论结果之前,让我们先讨论一下序列的提示(预填充)和生成(解码)部分之间的区别。正如我们将看到的,对于状态空间模型来说,预填充的细节比对于转换器模型更重要。当转换器生成下一个令牌时,它需要注意上下文中所有先前令牌的键和值。这意味着内存需求和生成时间随上下文长度进行线性缩放。状态空间模型仅关注并存储其循环状态,因此不需要额外的内存或时间来生成大型序列。虽然这解释了 SSM 在解码阶段相对于变压器的所谓优势,但预填充阶段需要额外的努力才能充分利用 SSM 架构。
预填充的标准方法是并行处理整个提示,以充分利用 GPU。这种方法在 optimum-benchmark 库中使用,我们将其称为并行预填充。并行预填充需要在内存中存储提示符中每个令牌的隐藏状态。对于变压器来说,这个额外的内存主要是由存储的KV缓存的内存主导的。对于 SSM 模型,不需要缓存,用于存储隐藏状态的内存成为唯一与提示长度成正比的组件。因此,内存需求将随着提示长度的增加而扩展,并且 SSM 模型将失去处理任意长序列的能力,类似于转换器。
并行预填充的替代方法是逐个令牌处理提示令牌,我们将其称为顺序预填充。与序列并行类似,它也可以在提示的较大块上完成,而不是在单个令牌上完成,以获得更好的 GPU 使用率。虽然顺序预填充对变压器来说意义不大,但它带回了 SSM 模型处理任意长提示的可能性。
考虑到这些评论,我们首先测试了单个 24 GB A10 GPU 上可以容纳的最大序列长度,并将结果如下图所示。批处理大小固定为 1,我们使用 float32 精度。即使对于并行预填充,Falcon Mamba 也可以适应比转换器更大的序列,而在顺序预填充中,它可以释放其全部潜力并可以处理任意长提示
接下来,我们使用批量大小 1 和 H100 GPU 在提示长度为 1 和最多 130k 生成的令牌的设置中测量生成吞吐量。结果如下图所示。我们观察到,我们的 Falcon Mamba 在恒定吞吐量下生成所有令牌,而 CUDA 峰值内存没有任何增加。对于 transformer 模型,峰值内存会增长,并且随着生成的 token 数量的增加,生成速度会减慢。
如何在 Hugging Face 变压器中使用它?
Falcon Mamba 架构将在 Hugging Face 转换器库 (>4.45.0) 的下一个版本中提供。要使用该模型,请确保安装最新版本的 Hugging Face 转换器或从源代码安装库。
Falcon Mamba 与您熟悉的大多数 Hugging Face 提供的 API 兼容,例如 AutoModelForCausalLM
或 pipeline
:
1 2 3 4 5 6 7 8 9 10 |
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 内存约束下运行模型,例如:
1 2 3 4 5 6 7 8 9 10 11 12 |
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) 数据令牌进行了微调。这种扩展训练增强了模型以更高的精度和有效性执行教学任务的能力。您可以通过我们的演示来体验 ininstruction 模型的功能,可在此处获得。对于聊天模板,我们使用以下格式:
1 2 3 |
<|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 负责变压器集成
- Vaibhav Srivastav、hysts 和 Omar Sanseviero 对与 Hub 相关的问题的支持
作者还要感谢 Tri Dao 和 Albert Gu 将 Mamba 架构实施并向社区开源。