純粋なCでのマンバモデルの推論
llama2.cからインスピレーションを受け、そのコードを使用しています
これはマンバSSMのリカレントモードのみを実装しています
関連するpytorchの実装と比較することができます
バッチには対応していません。コードは学習目的で最小限にしています。
それでも、CPU上のpytorchよりも速いです!!!
python3 tokenizer.py
python3 export.py state-spaces/mamba-130m model.bin
make fast
./mamba model.bin -n 20 -i "Customer Support should" -t 0.0
Pythonは、トークン化器とモデルをよりシンプルな形式にエクスポートするためにのみ使用されます(transformersとpytorchが必要です)
エクスポート部分で別のモデルを選択することができます
HuggingFaceに保存されているこれらのモデルを使用することができます:
state-spaces/mamba-130m
state-spaces/mamba-370m
state-spaces/mamba-790m
state-spaces/mamba-1.4b
state-spaces/mamba-2.8b
state-spaces/mamba-2.8b-slimpj
モデル名をexport.py
スクリプトの引数として指定することができます
エクスポートスクリプトは、モデルを(まだダウンロードされていない場合)hugingfaceのキャッシュディレクトリにダウンロードします。
オプションとして、手動でダウンロードした場合はモデルファイルへのパスも指定できます。例:
wget https://huggingface.co/state-spaces/mamba-130m/resolve/main/config.json?download=true -O config.json
wget https://huggingface.co/state-spaces/mamba-130m/resolve/main/pytorch_model.bin?download=true -O pytorch_model.bin
python3 export.py . model.bin
これはリカレントモデルであるため、内部状態を保存し、後でその状態に戻ることが可能です
内部状態のコピーを取得するには:
int state_size;
char* state = get_internal_state(mamba, &state_size);
内部状態を設定するには:
set_internal_state(mamba, state, state_size);
コードは3つのバージョンで利用可能で、それぞれ別のブランチにあります:
learning
- 基本的なものfused
- 基本的な機能をより大きなものに統合(比較することができます)cuda
- シンプルなGPU実装、理解しやすい
特殊文字に対しては、トークナイザーがさらに作業を必要とするかもしれません
自由に貢献し、PRを送ってください
MIT