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

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

🚀 快速阅读

1. 架构：基于 Mamba2 架构，旨在消除内存带宽瓶颈。
2. 性能：推理时吞吐量提升 2.5 倍，延迟加速 2 倍。
3. 应用：支持多平台，如 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 ..

        

          
        
        
        

          

          AI 代码解读
        
      
      

模型推理

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

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 代码解读
        
      
      

资源

• 项目官网：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 资讯和开源应用，也会不定期分享自己的想法和开源实例，欢迎关注我哦！

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