Bamba-9B-v2 - 快速而强大！
这是IBM、普林斯顿大学、卡内基梅隆大学和伊利诺伊大学香槟分校的联合研究成果。作者和贡献者的完整列表位于页面底部。
在2024年圣诞节期间，IBM、普林斯顿大学、卡内基梅隆大学和伊利诺伊大学香槟分校发布了Bamba v1，这是一个基于Mamba2的高性能预训练模型，具有完整的数据来源，训练量达到2T token。自那时起，我们一直忙于更新模型，加入了新的数据集。今天，我们很高兴发布Bamba v2，该模型额外训练了1T token，显著提升了Bamba v1的性能。L1和L2排行榜得分均优于Llama 3.1 8B，后者训练数据量是Bamba v2的近5倍。所有这些都得益于基于Mamba2架构带来的推理速度提升，最新vLLM的推理速度比同等大小的Transformer模型快2-2.5倍。
工件📦
快速而强大⚡🏎️
评估设置⚙️ 🖥️：我们按照此处的设置和脚本对所有模型重新运行了所有基准测试。对于v2排行榜结果，我们进行归一化并报告归一化后的结果。在所有评估中，除了特别指明的情况，分数越高越好。
虽然HF排行榜本身由于模型的演进已不再活跃，但比较关键基准是衡量模型能力的重要指标。我们将在下面提供各种基准的比较。我们观察到，与其他训练至少10T+ token（许多情况下甚至15T+）的SoTA模型相比，Bamba 9B v2在L1和L2平均得分上均优于流行的Llama 3.1 8B基础模型。
HF OpenLLM v1基准测试 + OpenbookQA, Boolq, 和 PIQA
| 模型 | 平均分 | MMLU | ARC-C | GSM8K | Hellaswag | OpenbookQA | Piqa | TruthfulQA | Winogrande | Boolq |
|---|---|---|---|---|---|---|---|---|---|---|
| Bamba 9B v2 | 62.63 | 67.92 | 63.57 | 41.70 | 83.85 | 51.0 | 83.62 | 50.86 | 79.48 | 82.78 |
| Nemotron-H 8B | 68.38 | 72.95 | 65.7 | 76.8 | 83.24 | 49.8 | 84.33 | 49.84 | 80.35 | 87.8 |
| Meta Llama 3.1 8B | 60.79 | 66.26 | 57.85 | 49.96 | 81.98 | 46.8 | 82.54 | 45.16 | 77.51 | 82.66 |
| Olmo2 7B | 63.99 | 63.96 | 64.51 | 68.01 | 81.93 | 49.2 | 81.39 | 43.32 | 770.3 | 84.77 |
| IBM Granite v3 8B | 64.25 | 64.13 | 63.74 | 60.2 | 83.34 | 47.2 | 83.08 | 51.35 | 79.79 | 87.22 |
| Gemma2 9B | 66.26 | 72.29 | 68.26 | 67.4 | 82.56 | 47.8 | 83.24 | 45.39 | 80.11 | 86.45 |
| Qwen2.5 7B | 69.05 | 75.41 | 63.82 | 83.24 | 80.23 | 48.4 | 81.28 | 56.34 | 75.93 | 87.74 |
HF OpenLLM v2基准测试
| 模型 | 平均分 | MMLU-PRO | BBH | GPQA | IFEval | MATH Lvl 5 | MuSR |
|---|---|---|---|---|---|---|---|
| Bamba 9B | 15.14 | 25.41 | 24.78 | 5.93 | 19.00 | 6.42 | 9.28 |
| Nemotron-H 8B | 23.98 | 36.02 | 32.91 | 10.96 | 25.15 | 23.56 | 15.3 |
| Meta Llama 3.1 8B | 14.45 | 25.46 | 25.16 | 8.61 | 12.55 | 6.19 | 8.72 |
| Olmo2 7B | 13.4 | 22.79 | 21.69 | 4.92 | 16.35 | 4.61 | 10.02 |
| IBM Granite v3 8B | 19.89 | 24.8 | 25.78 | 9.06 | 41.97 | 9.44 | 8.26 |
| Gemma2 9B | 21.79 | 34.84 | 34.81 | 11.07 | 21.28 | 13.44 | 15.3 |
| Qwen2.5 7B | 25.83 | 37.62 | 35.62 | 9.96 | 34.77 | 22.43 | 14.6 |
训练秘籍📖🍳
鉴于有限的GPU预算（192块A100），我们没有选择训练10T+ token的模型。相反，我们决定探索向现有模型中注入新数据，并尝试诸如模型合并等技术，这在Olmo2模型训练中有所提及。我们的训练秘籍如下图所示。
Bamba 9B v2的训练秘籍
我们使用了2T基础检查点（即Bamba 9b v1），并通过添加AllenAI作为Olmo2训练秘籍一部分发布的Olmo Mix来扩展它。我们使用恒定学习率调度，从2T到2.5T token，具体学习率为2e-5。精确的数据混合在下表中。
| 数据集 | 百分比 |
|---|---|
| DCLM | 90.5 |
| ArXiV | 1.14 |
| Pes2o | 2.36 |
| Starcoder | 4.5 |
| 代数堆栈 | 0.49 |
| 开放网络数学 | 0.51 |
| 维基 | 0.5 |
表1： 2-2.5T Olmo2混合数据
然后，我们使用了来自Nemotron-CC和Hugging Face数据集的混合合成数据，继续训练额外的500B token，使总训练量达到3T token。在此阶段，我们启动了两个作业（串行执行），对学习率调度进行消融实验，一个使用恒定学习率2e-5，另一个使用余弦学习率，从2e-5降至2e-6。我们的总体观察是，余弦和恒定学习率调度器是互补的。2.5T到3T的数据混合如下。
| 数据集 | 子数据集 | 百分比 |
|---|---|---|
| Nemotron-CC | CC | 70.5 |
| Nemotron-CC | 蒸馏 | 2.0 |
| Nemotron-CC | 问答对 | 5.0 |
| Nemotron-CC | 提取知识 | 2.0 |
| Nemotron-CC | 知识列表 | 0.25 |
| Nemotron-CC | Wrap medium | 0.25 |
| SmolLM | CosmopediaV2 | 2.0 |
| SmolLM | Fineweb edu dedup | 3.0 |
| Finemath | Finemath 3+ | 3.0 |
| Finemath | Infiwebmath 3+ | 3.75 |
| Dolmino | Pes2o | 2.5 |
| Dolmino | 维基 | 0.5 |
| Olmo | 代数堆栈 | 0.25 |
| Olmo | Starcoder | 5.0 |
表2： 2.5-3T自定义混合数据
最后，我们使用高质量数据对这两个模型进行退火，额外训练100B token，并使用MergeKit合并最终退火后的模型。具体来说，我们发现简单的加权平均对我们来说效果最好。由此产生的模型就是Bamba 9B v2！退火的数据混合如下。
| 数据集 | 子数据集 | 百分比 |
|---|---|---|
| Nemotron-CC | 蒸馏 | 7.0 |
| Nemotron-CC | 问答对 | 7.0 |
| Nemotron-CC | 提取知识 | 7.0 |
| Nemotron-CC | 知识列表 | 2.0 |
| Nemotron-CC | Wrap medium | 2.0 |
| Dolmino | DCLM | 20.0 |
| Dolmino | FLAN | 17.0 |
| Dolmino | Pes2o | 5.9 |
| Dolmino | 维基 | 3.7 |
| Dolmino | StackXchange | 1.3 |
| SmolLM | CosmopediaV2 | 17.5 |
| Finemath | Finemath 4+ | 9.6 |
表3： 3T-3.1T用于退火的自定义混合数据
关于指令调优的说明 🧠🔧📘 我们已经尝试使用Tuluv3数据来创建一个指令遵循模型，并利用Open Instruct。我们观察到，该模型在各项基准测试中的性能显著提高，L1平均值为64.69，L2平均值为24.68。我们正在开发一个具有非受限数据的指令遵循模型。
vLLM集成 📥🧠📤
我们正深入参与vLLM社区，以增加对Mamba2注意力的支持。在我们的首次发布中，我们与多个Mamba2模型开发者合作，增加了对张量并行支持并减少了内存复制。
我们目前正在进行三个重大更新：
(a) vLLM KV-缓存管理：页面的分配和KV-缓存的管理与基于Transformer的模型根本不同。我们正在与社区合作，以实现一个可以处理包括Mamba2在内的任何混合模型的通用解决方案。
(b) 分块预填充：我们深知分块预填充可以在实际工作负载中带来巨大的改进。虽然当前的内核对基于Transformer的模型表现良好，但我们需要为基于Mamba2的模型实现一个新的内核。
(c) 更快的解码内核：我们正在开发Conv1D内核的Triton实现，以确保模型在AMD GPU上也能高性能运行。此外，我们正在探索在解码时融合4-5个内核，以进一步减少token生成的延迟。
我们预计，通过上述更改，我们可以为基于Mamba2的模型实现一流的支持，并显著优于目前的性能（对于较长序列长度/较大批量，可达到2-2.5倍）。
行动呼吁📢👉🚀
我们致力于提供开放数据集，并完整复现我们的结果。我们呼吁社区在多个方面帮助改进模型：
- 模型上的测试时间扩展和GRPO
- 提高vLLM的推理性能（我们期望比相应的Transformer模型快4-5倍）
- 帮助我们提高数学能力！
贡献者
- 数据收集和整理：我们感谢并鸣谢AllenAI团队提供了高质量的开源数据集Dolma、OlmoMix和DolminoMix，Hugging Face数据团队提供了FineWeb*和SmolLM语料库，以及NVIDIA团队创建了Nemotron-CC。这些都是巨大的贡献，使我们能够创建和改进模型。
- 数据预处理：我们感谢IBM内部数据预处理团队，特别是Tuan Hoang Trong、Syed Zawad、Jay Gala和Ryan Gordon在大规模数据标记方面的帮助。标记代码可在此处获取。
- 模型架构：模型架构设计由普林斯顿大学、卡内基梅隆大学、IBM和伊利诺伊大学香槟分校联合完成，参与人员包括：Tri Dao（普林斯顿大学）、Albert Gu（卡内基梅隆大学）、Linsong Chu（IBM）、Davis Wertheimer（IBM）、Minjia Zhang（伊利诺伊大学香槟分校）、Mudhakar Srivatsa（IBM）和Raghu Ganti（IBM）。
- 模型训练：模型训练主要由IBM团队完成，使用了Tri Dao和Albert Gu的Mamba2内核和层实现。IBM团队的主要参与人员包括：Linsong Chu、Divya Kumari、Davis Wertheimer、Garrett Goon、Raghu Ganti和Dakshi Agrawal。
- 模型调优：模型的调优由IBM团队在TRL中启用和验证，参与人员包括Sukriti Sharma和Anh Uong。
- 模型推理：
transformers、vLLM和llama.cpp中的模型推理基于普林斯顿大学和卡内基梅隆大学编写的内核。IBM团队正与社区合作，使其能在各种生态系统中运行。团队成员包括：Fabian Lim、Antoni viros i Martin、Adnan Hoque、Jamie Yang、Nelson Nimura Gonzalez、Tuan Hoang Trong、Joshua Rosenkranz、Nick Hill和Gabe Goodhart。 - 量化：量化由IBM团队负责——王乃刚和刘驰淳。
- 评估：评估由IBM团队主导，长上下文评估由伊利诺伊大学香槟分校执行，参与人员包括：Yotam Perlitz、Ofir Arviv、Michal Shmueli-Scheuer（IBM）、Haoechen Shen和Minjia Zhang（伊利诺伊大学香槟分校）。
最后，我们要感谢我们的领导在这项工作中的支持——Priya Nagpurkar、David Cox、Sriram Raghavan、Aya Soffer、Ruchir Puri和Mukesh Khare。
我们还要感谢社区——Hugging Face在transformer方面的支持以及vLLM在推理优化方面的贡献。
非常感谢Meta PyTorch、AllenAI和Hugging Face团队对开放计划的贡献，PyTorch FSDP让我们顺利训练了这款模型，并且各种高质量的数据集使我们能够进一步更新模型！
