mamba-2.8b-hf

mamba-2.8b-hf项目介绍

mamba-2.8b-hf项目是一个与transformers库兼容的模型库，提供了一个完整的config.json配置文件和分词器。这些模型检查点在项目中保持原样。

安装说明

要使用mamba-2.8b-hf模型，用户需要首先从main分支安装transformers库，直到transformers=4.39.0版本发布。安装命令如下：

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

此外，推荐同时安装两个依赖库：causal_conv_1d和mamba-ssm。它们的安装命令如下：

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

如果未安装这两个库，将会使用“eager”实现；如果安装了，则使用更优化的cuda内核。

文本生成

用户可以使用传统的generate API来进行文本生成，具体示例如下：

from transformers import MambaConfig, MambaForCausalLM, AutoTokenizer
import torch

tokenizer = AutoTokenizer.from_pretrained("state-spaces/mamba-2.8b-hf")
model = MambaForCausalLM.from_pretrained("state-spaces/mamba-2.8b-hf")
input_ids = tokenizer("Hey how are you doing?", return_tensors="pt")["input_ids"]

out = model.generate(input_ids, max_new_tokens=10)
print(tokenizer.batch_decode(out))
# 输出: ["Hey how are you doing?\n\nI'm doing great.\n\nI"]

使用PEFT进行微调

mamba-2.8b-hf模型可以通过peft库进行微调，建议将模型保持在float32格式。下面是一个示例代码：

from datasets import load_dataset
from trl import SFTTrainer
from peft import LoraConfig
from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments

tokenizer = AutoTokenizer.from_pretrained("state-spaces/mamba-2.8b-hf")
model = AutoModelForCausalLM.from_pretrained("state-spaces/mamba-2.8b-hf")
dataset = load_dataset("Abirate/english_quotes", split="train")

training_args = TrainingArguments(
    output_dir="./results",
    num_train_epochs=3,
    per_device_train_batch_size=4,
    logging_dir='./logs',
    logging_steps=10,
    learning_rate=2e-3
)

lora_config = LoraConfig(
    r=8,
    target_modules=["x_proj", "embeddings", "in_proj", "out_proj"],
    task_type="CAUSAL_LM",
    bias="none"
)

trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    args=training_args,
    peft_config=lora_config,
    train_dataset=dataset,
    dataset_text_field="quote",
)

trainer.train()