Mamba 7b Rw

🚀 Mamba-7B

Mamba-7B是一个具有70亿参数的模型，采用Mamba架构，在RefinedWeb数据集上进行了多轮训练（处理了1.2万亿个标记）。Mamba是一种状态空间模型，与标准的Transformer架构不同，它不使用自注意力机制。该模型在各种自然语言基准测试中表现出色。截至目前，公开发布的最大纯Mamba预训练模型是Mamba-2.8B。本项目遵循其训练方案，发布了Mamba-7B的版本。此模型作为论文Linearizing Large Language Models的基线模型进行训练。

🚀 快速开始

本模型使用OpenLM进行训练，权重已转换为与HuggingFace兼容的格式。以下是使用示例：

from transformers import AutoTokenizer, AutoModelForCausalLM
tokenizer = AutoTokenizer.from_pretrained("tri-ml/mamba-7b-rw")
model = AutoModelForCausalLM.from_pretrained("tri-ml/mamba-7b-rw")

inputs = tokenizer(["The Toyota Supra"], return_tensors="pt")
gen_kwargs = {"max_new_tokens": 50, "top_p": 0.8, "temperature": 0.8, "do_sample": True, "repetition_penalty": 1.1}
output = model.generate(inputs['input_ids'], **gen_kwargs)
output = tokenizer.decode(output[0].tolist(), skip_special_tokens=True)
print(output)
# The Toyota Supra is a sports car that has been in production since 1978. The car was discontinued in 2002, but it has recently been revived and will be available again in 2020. The Supra has always been known for its powerful engines and agile handling. 

✨ 主要特性

• 采用Mamba架构，不使用自注意力机制，在自然语言基准测试中表现出色。
• 作为论文Linearizing Large Language Models的基线模型进行训练。

📦 安装指南

文档未提及具体安装步骤，可参考OpenLM的相关说明。

💻 使用示例

基础用法

from transformers import AutoTokenizer, AutoModelForCausalLM
tokenizer = AutoTokenizer.from_pretrained("tri-ml/mamba-7b-rw")
model = AutoModelForCausalLM.from_pretrained("tri-ml/mamba-7b-rw")

inputs = tokenizer(["The Toyota Supra"], return_tensors="pt")
gen_kwargs = {"max_new_tokens": 50, "top_p": 0.8, "temperature": 0.8, "do_sample": True, "repetition_penalty": 1.1}
output = model.generate(inputs['input_ids'], **gen_kwargs)
output = tokenizer.decode(output[0].tolist(), skip_special_tokens=True)
print(output)
# The Toyota Supra is a sports car that has been in production since 1978. The car was discontinued in 2002, but it has recently been revived and will be available again in 2020. The Supra has always been known for its powerful engines and agile handling. 

📚 详细文档

模型详情

训练详情

• Mamba-7B在128个H100 80GB GPU上使用AWS SageMaker进行训练。
• 训练于2024年3月开始，持续了三周。 | 超参数 | 值 | |------------|---------| | 精度 | bfloat16 | | 优化器 | AdamW | | 学习率 | 3e-4 | | 学习率冷却结束值 | 1e-5 | | 预热步数 | 2000 | | Z损失 | 1e-4 | | 批量大小 | 200万 |

性能评估

评估使用Eleuther LM Eval Harness仓库进行。以下是Mamba 7B与其他基础模型的性能对比：

🔧 技术细节

• Mamba是一种状态空间模型，不使用自注意力机制，在自然语言处理任务中表现良好。
• 模型训练使用了AWS SageMaker和特定的超参数设置，以确保模型的性能和稳定性。

📄 许可证

本模型遵循Apache License, Version 2.0许可协议。

如何引用

如果使用此模型，请引用论文Linearizing Large Language Models：

@article{Mercat2024Linearizing,
  title={Linearizing Large Language Models},
  author={Jean Mercat and Igor Vasiljevic and Sedrick Keh and Kushal Arora and Achal Dave and Adrien Gaidon and Thomas Kollar},
  journal={arXiv preprint arXiv:2405.06640},
  year={2024}
}

引用文献

Mamba

@article{mamba,
  title={Mamba: Linear-Time Sequence Modeling with Selective State Spaces},
  author={Gu, Albert and Dao, Tri},
  journal={arXiv preprint arXiv:2312.00752},
  year={2023}
}

OpenLM

@misc{open_lm,
  author = {Gururangan, Suchin and Wortsman, Mitchell and Gadre, Samir Yitzhak and Dave, Achal and Kilian, Maciej and Shi, Weijia and Mercat, Jean and Smyrnis, Georgios and Ilharco, Gabriel and Jordan, Matt and Heckel, Reinhard and Dimakis, Alex and Farhadi, Ali and Shankar, Vaishaal and Schmidt, Ludwig},
  title = {{open_lm}:  a minimal but performative language modeling (LM) repository},
  year = {2023},
  note = {GitHub repository},
  url = {https://github.com/mlfoundations/open_lm/}
}