Вывод моделей Mamba на чистом C
Вдохновлен и использует код из llama2.c
Это реализует только рекуррентный режим Mamba SSM
Вы можете сравнить его с соответствующей реализацией на pytorch
Нет поддержки для пакетов. Код минимален для обучения.
Тем не менее, он быстрее, чем pytorch на CPU!!!
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