Axon Kit

Kit — лёгкая и быстрая модель системы Axon компании SableAI. Это её первая выпущенная версия: небольшая русскоязычная языковая модель (GPT, 89,3 млн параметров), обученная с нуля на одной домашней видеокарте с 6 ГБ памяти. Веса не взяты ни у какой другой модели.

Это базовая модель, а не чат-бот. Она продолжает текст, а не отвечает на вопросы. Она маленькая и обучена на небольшом корпусе, поэтому часто ошибается в фактах и теряет мысль. Не используйте её там, где нужна надёжность.

В transformers и vLLM она не загрузится: архитектура собственная. Код модели лежит рядом (model.py), запуск описан ниже.

О проекте

SableAI (sable — соболь) — компания, Axon — её ИИ-система (нейрон, связь, инженерия). Модели системы Axon названы по животным семейства куньих и выстроены от простого к сложному:

Модель Для чего Статус
Kit простые и мгновенные задачи: лёгкая и быстрая, как детёныш соболя первая версия выпущена (89,3 млн параметров, эта страница)
Stout офисная работа и программирование: крепкая, надёжная, рабочая в планах
Marten сложные аналитические задачи: умная и хитрая, как куница в планах
Fisher глубокий научный и творческий потенциал, флагман в планах

У Kit будут версии разного размера и с разными языками. Эта версия — стартовая точка, на которой отлаживаются данные, токенайзер и обучение. Дальнейшие планы: английский язык наравне с русским, затем другие языки и код. Всё это пока не реализовано, и у выпущенной модели этих возможностей нет.

Характеристики

параметры 89,3 млн (веса токенов и выходного слоя общие)
слои / голова / размерность 16 / 10 / 640
контекст 512 токенов
словарь 16 000, byte-level BPE, обучен на том же корпусе
архитектура decoder-only трансформер, LayerNorm, GELU, обучаемые позиции
формат весов model.safetensors, fp32 (полная точность, без сжатия), около 357 МБ

Запуск

pip install torch tokenizers safetensors huggingface_hub
hf download SableAI-Dev/Axon-Kit --local-dir Axon-Kit
cd Axon-Kit
python generate.py model.safetensors "Москва — это" --tokens 150

Параметры генерации: --temp (по умолчанию 0.8), --top_k (40). Нужна видеокарта с CUDA. Один запуск занимает меньше 1 ГБ видеопамяти.

Данные

Русский срез FineWeb-2 (rus_Cyrl): документы короче 300 символов отброшены, оставлено около 131M токенов. Последний 1% корпуса не участвует в обучении и используется для проверки. Дедупликации и дополнительной фильтрации сверх этого не было. Корпус собран из веб-страниц, поэтому модель может воспроизводить их смещения и ошибки.

Обучение

железо одна NVIDIA RTX 2060 (6 ГБ), Intel Core i5-12400F, 32 ГБ ОЗУ
оптимизатор AdamW, betas 0.9/0.95, weight decay 0.1 (только для матриц)
скорость обучения 3e-4, прогрев 200 шагов, затем косинус до 3e-5
батч 64 последовательности по 512 токенов (8 × накопление 8), 32 768 токенов на шаг
точность fp16 с автоматическим масштабированием потерь, обрезка градиента 1.0
скорость около 23 500 токенов в секунду
этапы 14 000 шагов с батчем 8 (57M токенов); 1 200 шагов с батчем 64 (39M токенов); ещё 5 000 шагов с батчем 64 (164M токенов, около 2 часов). Всего примерно 260M токенов, около двух проходов по корпусу

Чтобы не упираться в видеокарту, весь корпус держится в оперативной памяти, а батчи готовятся на процессоре в фоновых потоках, пока видеокарта считает.

Качество

Val loss на отложенном 1% корпуса на последнем этапе обучения (5 000 шагов):

шаг 250 1000 2000 3000 4000 5000
val loss 3.812 3.698 3.629 3.556 3.494 3.483

Итоговый val loss 3.48 (перплексия около 33 на токенах словаря 16k). Оценка идёт по 40 батчам, поэтому небольшие колебания между точками (например, 3.478 на шаге 4750) — шум, а не тренд. Другие бенчмарки не запускались.

Примеры продолжения текста (температура 0.8, top-k 40):

Москва — это не «прибытие» и в самом деле не что иное, как смерть. Для нас это не идеальная задача, — в ней много нечистоплюдая и безразличная, где «высокая» и безотносимая, даже «абсолютная»

Для того чтобы приступить к работе, приступить к работе, а после этого вы сможете научиться учить свою деятельность.

Текст грамматически похож на русский и по стилю на веб-страницы, но по смыслу бессвязен: повторы, выдуманные слова и ошибки в согласовании встречаются постоянно.

Ограничения

  • Только русский язык; на других языках модель почти бесполезна.
  • Контекст 512 токенов.
  • Нет дообучения на инструкциях и диалогах, нет фильтров безопасности.
  • Обучена на ограниченном объёме данных (меньше нескольких эпох по 131M токенов), поэтому заметно недообучена.

Лицензия

Код и веса под Apache 2.0. Обучающие данные собраны из FineWeb-2 (ODC-By 1.0) и сохраняют исходные условия.

Downloads last month
494
Safetensors
Model size
89.3M params
Tensor type
F32
·
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Dataset used to train SableAI-Dev/Axon-Kit