Bamba-9B：基于 Mamba2 架构的仅解码语言模型，旨在提高大型语言模型在推理时的效率

蚝油菜花
109 阅读3分钟

❤️ 如果你也关注 AI 的发展现状，且对 AI 应用开发非常感兴趣，我会每日跟你分享最新的 AI 资讯和开源应用，也会不定期分享自己的想法和开源实例，欢迎关注我哦！

🥦 微信公众号｜搜一搜：蚝油菜花 🥦

🚀 快速阅读

  1. 架构：基于 Mamba2 架构，旨在消除内存带宽瓶颈。
  2. 性能：推理时吞吐量提升 2.5 倍，延迟加速 2 倍。
  3. 应用：支持多平台，如 transformers、vLLM、TRL 和 llama.cpp。

正文（附运行示例）

Bamba-9B 是什么

公众号: 蚝油菜花 - bamba

Bamba-9B 是由 IBM、普林斯顿大学、卡内基梅隆大学和伊利诺伊大学香槟分校联合推出的基于 Mamba2 架构的仅解码语言模型。该模型在完全开放的数据集上训练，旨在提高大型语言模型的推理效率，特别是在处理长文本时的内存带宽瓶颈。

Bamba-9B 在推理时相较于标准变换器模型展现出 2.5 倍的吞吐量提升和 2 倍的延迟加速。模型训练使用 2.2 万亿个 token，进一步验证了新兴架构的潜力，在保持与最先进的变换器模型竞争的同时，提供更高的推理效率。

Bamba-9B 的主要功能

  • 提高推理效率：Bamba-9B 设计的主要目标是提高大型语言模型在推理时的效率，特别是在处理长文本时的内存带宽瓶颈。
  • 吞吐量和延迟优化：相比于标准变换器模型，Bamba-9B 在推理时展示 2.5 倍的吞吐量提升和 2 倍的延迟加速。
  • 开放数据集训练：Bamba-9B 完全在开放数据集上训练，有助于社区进行透明度和可复制性的实验。
  • 多平台支持：Bamba-9B 支持在多个开源平台使用，如 transformers、vLLM、TRL 和 llama.cpp。

Bamba-9B 的技术原理

  • 混合 Mamba2 架构：基于 Mamba2 架构，一种新兴的架构，使 KV-cache 大小恒定，消除内存带宽瓶颈。
  • 恒定 KV-cache：KV-cache 所需的内存量随上下文长度增加而增加，而 Mamba2 架构保持 KV-cache 大小不变，解决这一问题。
  • 两阶段训练方法：采用两阶段训练方法，第一阶段用 Dolma v1.7 数据集进行训练，第二阶段用 Fineweb-edu 和 Cosmopedia 等高质量数据集进行额外训练。
  • 数据加载器：推出一个分布式状态无关的数据加载器，支持大规模分布式训练，并与 Torch Titan 集成。
  • 量化技术：支持模型量化，基于 llm-compressor 将模型量化到 fp8，减少模型大小，提高推理速度，同时保持准确性。
  • 上下文长度扩展：Bamba-9B 正在探索长上下文长度扩展的方法，如将 LongRope 应用于全注意力层，处理更长的上下文。

如何运行 Bamba-9B

安装依赖

除了 PyTorch，还需要安装一些额外的依赖项：

git clone https://github.com/Dao-AILab/causal-conv1d.git
cd causal-conv1d && pip install . && cd ..
git clone https://github.com/state-spaces/mamba.git
cd mamba && pip install . && cd ..
git clone https://github.com/Dao-AILab/flash-attention.git
cd flash-attention && pip install . && cd ..

模型推理

使用以下命令进行文本生成：

python text_generation.py \
--model_path ibm-fms/Bamba-9B \
--tokenizer_path ibm-fms/Bamba-9B \
--prompt "The largest living mammal on Earth is " \
--max_new_tokens 128

资源

❤️ 如果你也关注 AI 的发展现状，且对 AI 应用开发非常感兴趣，我会每日跟你分享最新的 AI 资讯和开源应用，也会不定期分享自己的想法和开源实例，欢迎关注我哦！

🥦 微信公众号｜搜一搜：蚝油菜花 🥦

avatar
文章
阅读
粉丝