AliceAI-Foundation

English version

AliceAI-Foundation-80B-A3B-Base – базовая языковая модель с гибридной архитектурой и MoE-слоями. Модель содержит 80 млрд параметров, из которых для каждого токена активируются 3 млрд, и поддерживает контекст длиной до 262 144 токенов. Модель была обучена полностью с нуля.

При создании модели мы заново собрали обучающий корпус, выбрали архитектуру и гиперпараметры, а также подготовили данные для сложных рассуждений и взаимодействия с инструментами. Ключевые решения проверялись в серии отдельных обучений с нуля объёмом по 2 трлн токенов каждое.

В математике, программировании и других задачах на рассуждения модель показывает результаты на уровне более крупных опенсорс-моделей, а особенно сильна в задачах на фактические знания на русском языке. Вместе с весами мы публикуем фактологические бенчмарки WikiWebFacts и HardMultiQA, ориентированные на русскоязычный контекст, и протоколы их оценки.

Подробно про эту модель можно прочитать в статье на Хабре.

Обзор модели

• Тип: авторегрессионная языковая модель
• Этап обучения: предобучение
• Языковая модель
  • Количество параметров: 80B всего, 3B активных
  • Размер скрытого состояния: 2048
  • Размер словаря: 129024
  • Количество слоев: 48
  • Схема слоёв: 12 × (3 × (KDA → MoE) → 1 × (Gated Attention → MoE))
  • KDA:
    • Количество query-голов: 32
    • Количество KV-голов: 32
    • Размерность query-головы: 128
    • Размерность KV-головы: 128
    • Рамер ядра свертки: 4
  • Gated Attention:
    • Количество query-голов: 16
    • Количество KV-голов: 2
    • Размерность query-головы: 256
  • MoE:
    • Количество экспертов: 512
    • Top-K: 10 + 1 общий эксперт
    • Промежуточная размерность эксперта: 512
  • MTP: 1 слой
• Длина контекста: 262144

Бенчмарки

Как использовать

Transformers

Модель можно запустить через Transformers. Референсная версия Transformers — 5.16.1. Для выполнения KDA-слоёв на GPU требуется flash-linear-attention с поддержкой KDA:

pip install \
  transformers==5.16.1 \
  accelerate==1.14.0 \
  flash-linear-attention==0.5.0

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "yandex/AliceAI-Foundation-80B-A3B-Base"

tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    trust_remote_code=True,
    dtype=torch.bfloat16,
    device_map="auto",
)

prompt = "Есть 256 монет с разным весом, за какое минимальное количество попарных взвешиваний можно найти вторую по весу монету?"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
output_ids = model.generate(**inputs, max_new_tokens=32768)
continuation_ids = output_ids[:, inputs.input_ids.shape[1] :]
print(tokenizer.decode(continuation_ids[0], skip_special_tokens=True))

vLLM

Также модель можно запустить через vLLM. Для запуска требуются Docker и NVIDIA Container Toolkit.

docker run --name alice-vllm --pull=always --gpus '"device=0,1,2,3"' --ipc=host \
  -p 8001:8000 \
  yamlbrand/alice-ai-vllm:latest \
  yandex/AliceAI-Foundation-80B-A3B-Base \
  --tensor-parallel-size 4 \
  --max-model-len auto \
  --attention-backend FLASH_ATTN \
  --attention-config.flash_attn_version=2 \
  --speculative-config '{"method":"mtp","num_speculative_tokens":1}'

Повторный запуск после остановки, с сохранённым кешем:

docker start -a alice-vllm

Чтобы использовать все GPU, замените --gpus '"device=0,1,2,3"' на --gpus all и укажите соответствующее значение размера тензорного параллелизма.

После запуска сервера отправьте запрос:

curl http://127.0.0.1:8001/v1/completions \
  -H 'Content-Type: application/json' \
  -d '{
    "model": "yandex/AliceAI-Foundation-80B-A3B-Base",
    "prompt": "Есть 256 монет с разным весом, за какое минимальное количество попарных взвешиваний можно найти вторую по весу монету?",
    "max_tokens": 32768,
    "temperature": 0
  }'

Токенизатор

Токенизатор загружается как LlamaTokenizer из файла tokenizer.model. Модель токенизации использует SentencePiece BPE. Маркеры [COT_ENABLE], [COT_START], [COT_END], а также маркеры инструментов являются обычными атомарными токенами словаря, а не специальными токенами Hugging Face.

В tokenizer_config.json для параметра legacy явно установлено значение false, чтобы сохранить ожидаемую обработку пробелов. Не переопределяйте его значением true.

Как дообучить под свои задачи

Формат данных

Для подготовки агентских данных, использованных при обучении модели, мы использовали стандартный формат OpenAI Messages. В нём траектория задаётся последовательностью сообщений с ролями system, user, assistant, tool и meta, а определения доступных инструментов передаются отдельно в поле tools.

Перед токенизацией каждая такая траектория рендерилась с помощью chat_template.jinja. Шаблон задаёт префиксы ролей, представление reasoning-трейсов, описаний инструментов, вызовов функций и результатов их выполнения. Именно в таком текстовом представлении эти данные использовались при обучении модели.

Для sft и RL рекомендуем хранить данные в формате OpenAI Messages и рендерить их с помощью этого шаблона. Так формат новых данных будет совпадать с форматом, который модель видела во время претрейна.

Мы намеренно не указываем этот шаблон как chat_template в tokenizer_config.json: Alice-AI-Foundation-80B-A3B-Base — базовая модель, поэтому у неё нет единственного формата диалога, который должен автоматически применяться при инференсе. Приведённый шаблон предназначен именно для подготовки данных к дообучению.

Пример LoRA-дообучения

В репозитории есть минимальный пример PEFT-дообучения finetune_lora.py. Он загружает зафиксирова