Mamba环境配置

 环境说明

Requirements		本环境
Linux		ubuntu 20.04
NVIDIA GPU	RTX 20+	RTX A6000
CUDA	11.6+	11.8
PyTorch	1.12.0+	2.1.1
gcc	9.0.0+	9.4.0
causal-conv1d	1.2.0+	1.2.0
mamba	/	1.2.0
Python	/	3.9.12

 Pre、查看gcc版本

gcc -v

保证gcc版本不低于9.0

一、 检查并更新显卡驱动

1.1. 查看当前显卡驱动版本

nvidia-smi -a

 目前的显卡驱动是470.57.02，CUDA Driver是11.4，显卡是A6000

1.2. 更新显卡驱动 (CUDA Driver.version >= CUDA Toolkit.version ==11.8)

NVIDIA官网查看目前的最新驱动的版本

在宿主机中打开ubuntu软件更新器，选择最新的版本

重启

sudo reboot

 显示显卡驱动已更新

二、容器内创建虚拟环境mamba

 3.1.创建conda虚拟环境

python --version                    # 查看当前python版本
conda create -n mamba python=3.8.5  # 创建虚拟环境, 名为mamba 
# conda remove -n xx --all          # 删除环境，防止手误创建错误
conda activate mamba # 激活环境

3.2. 安装cudatoolkit

conda install cudatoolkit==11.8 -c nvidia
conda install -c "nvidia/label/cuda-11.8.0" cuda-nvcc

 3.3. 安装pytorch和torchvision

pip install torch==2.1.1 torchvision==0.16.1 torchaudio==2.1.1 --index-url https://download.pytorch.org/whl/cu118
conda install packaging

3.4. 安装Transformers

pip install transformers
pip install einops

3.4. 安装casual-conv1d和mamba-ssm

注意casual-conv1d和mamba-ssm的版本一定要对应上，否则报错。

cd dir # 进入操作的目录
git clone https://github.com/Dao-AILab/causal-conv1d.git
cd causal-conv1d
git checkout v1.2.0     # current latest version tag
python setup.py install

cd dir # 进入操作的目录
git clone https://github.com/state-spaces/mamba.git
cd mamba
git checkout v1.2.0 # current latest version tag
python setup.py install

四、Demo推理

下载mamba-130m的权重文件https://huggingface.co/state-spaces/mamba-130m-hf/tree/main

将全部文件下载到名为mamba-130m-hf的文件夹中

新建一个demo.py文件

from transformers import MambaConfig, MambaForCausalLM, AutoTokenizer
import torch
path="./path/to/mamba-130m-hf" # 修改这里的路径
# 加载模型
tokenizer = AutoTokenizer.from_pretrained(path,local_files_only=True)
model = MambaForCausalLM.from_pretrained(path,local_files_only=True)
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))

运行

五、微调

 Practical ML Dive - How to train Mamba for Question Answering