Mambaoutai

🚀 Mambaoutai 1.6B

Mambaoutai是此博客文章中描述的所有实验和训练运行的成果，该文章分享了该模型系列的所有细节。Mambaoutai是一系列小型Mamba检查点，供社区探索使用，在法语、英语和代码数据上进行训练。我们使用WSD调度器运行了两个不同的衰减阶段，并发布了有无指令数据预训练的模型检查点。

🚀 快速开始

Mambaoutai可用于文本生成、推理等任务。你可以按照以下步骤使用该模型。

✨ 主要特性

• 多语言支持：在法语、英语和代码数据上进行训练。
• 不同训练阶段检查点：发布了有无指令数据预训练的模型检查点。
• 轻量级模型：仅有1.6B参数，可在CPU上以合理速度运行。

📦 安装指南

你需要从main分支安装transformers，直到transformers=4.39.0版本发布。

pip install git+https://github.com/huggingface/transformers@main

我们还建议你使用以下命令安装causal-conv1d和mamba-ssm：

pip install causal-conv1d>=1.2.0
pip install mamba-ssm>=1.2.0

如果这两个库未安装，将使用“eager”实现（不推荐），否则将使用更优化的CUDA内核。

💻 使用示例

基础用法

使用以下代码片段从模型生成文本：

from transformers import MambaConfig, MambaForCausalLM, AutoTokenizer
import torch

if model_has_instruct_data:
    # 使用聊天令牌
    prompt = ”<start_user>Tell me something about Paris.<end_message><start_assistant>”
else:
    # 温和地提示未经过指令调优的模型
    prompt = ”This is a text about Paris. Paris is”

tokenizer = AutoTokenizer.from_pretrained("lightonai/mambaoutai")
model = MambaForCausalLM.from_pretrained("lightonai/mambaoutai")
input_ids = tokenizer(prompt, return_tensors="pt")["input_ids"]

out = model.generate(input_ids, max_new_tokens=10)
print(tokenizer.batch_decode(out))

高级用法

你可以在仓库分支中找到一些训练检查点。在训练过程中的某个时间点对应的分支上。你可以通过在from_pretrained方法中添加revision参数，使用这些训练检查点进行推理。例如，要加载预训练30000步后的模型检查点，可以使用以下代码：

from transformers import MambaConfig, MambaForCausalLM, AutoTokenizer
import torch

tokenizer = AutoTokenizer.from_pretrained("lightonai/mambaoutai", revision="pre-30000")
model = MambaForCausalLM.from_pretrained("lightonai/mambaoutai", revision="pre-30000")
input_ids = tokenizer("What is a mamba?", return_tensors="pt")["input_ids"]

out = model.generate(input_ids, max_new_tokens=10)
print(tokenizer.batch_decode(out))

设备上推理

由于Mambaoutai仅有1.6B参数，它可以在CPU上以合理速度运行。以下是在llama.cpp上运行它的示例：

# 克隆llama.cpp仓库并从源代码编译
git clone https://github.com/ggerganov/llama.cpp
cd llama.cpp
make

# 创建虚拟环境并安装依赖
conda create -n mamba-cpp python=3.10
conda activate mamba-cpp
pip install -r requirements/requirements-convert-hf-to-gguf.txt

# 从本仓库下载权重、分词器、配置、分词器配置和特殊令牌映射，并将它们放在目录'Mambaoutai/'中
mkdir Mambaoutai

# 将权重转换为GGUF格式
python convert-hf-to-gguf.py Mambaoutai

# 使用提示进行推理
./main -m Mambaoutai/ggml-model-f16.gguf -p "Building a website can be done in 10 simple steps:\nStep 1:" -n 400 -e -ngl 1

🔧 技术细节

训练硬件

无指令数据的模型检查点在OVH Cloud提供的NVIDIA DGX H100上进行了全面训练，而有指令数据的衰减阶段在Orange Cloud Avenue的HPE Cray（配备8xH100）上进行。消融实验在MeluXina的16个节点（4xA100 - 40GB）上进行。

模型超参数

模型超参数的更多细节如下表所示：

📄 许可证

本项目采用Apache-2.0许可证。

数据集

• togethercomputer/RedPajama-Data-V2
• stingning/ultrachat

语言

• fr
• en

评估指标

• accuracy
• perplexity