❤️ 如果你也关注 AI 的发展现状，且对 AI 应用开发非常感兴趣，我会每日跟你分享最新的 AI 资讯和开源应用，也会不定期分享自己的想法和开源实例，欢迎关注我哦！
🥦 微信公众号｜搜一搜：蚝油菜花 🥦
🚀 快速阅读
- 架构：基于 Mamba2 架构，旨在消除内存带宽瓶颈。
- 性能：推理时吞吐量提升 2.5 倍，延迟加速 2 倍。
- 应用：支持多平台，如 transformers、vLLM、TRL 和 llama.cpp。
正文（附运行示例）
Bamba-9B 是什么
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
资源
- 项目官网：https://huggingface.co/collections/ibm-fms/bamba
- GitHub 仓库：https://github.com/foundation-model-stack/bamba
- HuggingFace 模型库：https://huggingface.co/collections/ibm-fms/bamba-674f1388b9bbc98b413c7bab
- arXiv 技术论文：https://arxiv.org/pdf/2410.06734
❤️ 如果你也关注 AI 的发展现状，且对 AI 应用开发非常感兴趣，我会每日跟你分享最新的 AI 资讯和开源应用，也会不定期分享自己的想法和开源实例，欢迎关注我哦！
🥦 微信公众号｜搜一搜：蚝油菜花 🥦
