mamba-370m-hf - 兼容transformers库的高效语言模型

项目简介

mamba-370m-hf 项目是一个与 transformers 库兼容的模型，实际上是 mamba-2.8b 的检查点版本。该库保存了完整的 config.json 和分词器文件，这些都是模型运行所必需的。

使用指南

在使用 mamba-370m-hf 项目之前，用户需要从 transformers 的 main 分支安装相关库，直到 transformers=4.39.0 版本发布为止。安装命令如下：

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

此外，建议安装 causal_conv_1d 和 mamba-ssm 两个库，以便使用优化的 cuda 内核。如果这两个库没有安装，系统将默认使用一个“eager”的实现。

安装命令如下：

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

文本生成

可以使用经典的 generate API 来生成文本。以下是一个简单的代码示例：

from transformers import MambaConfig, MambaForCausalLM, AutoTokenizer
import torch

tokenizer = AutoTokenizer.from_pretrained("state-spaces/mamba-370m-hf")
model = MambaForCausalLM.from_pretrained("state-spaces/mamba-370m-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 微调示例

如果想使用 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-370m-hf")
model = AutoModelForCausalLM.from_pretrained("state-spaces/mamba-370m-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()

以上示例展示了如何加载数据集、配置训练参数、以及利用 LoraConfig 进行微调配置，并最终开始训练。这种方法便于针对具体文本任务进行模型优化。