AliceAI-Foundation-80B-A3B-Base

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

Бенчмарки

Названия русскоязычных бенчмарков выделены зелёным, англоязычных — синим.

Все замеры в этом разделе проведены во внутренней инфраструктуре замеров, инференс в фреймворке vllm с t=0 для всех моделей. Жирным выделен победитель в каждой строчке.

Все замеры в этом разделе проведены во внутренней инфраструктуре замеров, инференс в фреймворке vllm с t=1 и штрафами за повторы (repetition_penalty=1, presence_penalty=1.5) для всех моделей. Жирным выделен победитель в каждой строчке.

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

Transformers

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

python3 -m venv .venv
source .venv/bin/activate
pip install \
  transformers[sentencepiece]==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.rstrip(), 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

Доступны два способа запуска:

• Официальная nightly-сборка vLLM без MTP.
• Готовый Docker-образ с поддержкой MTP.

Стандартный запуск

Установите nightly-сборку из ветки main в чистое окружение:

uv pip install -U vllm --torch-backend=auto --extra-index-url https://wheels.vllm.ai/nightly

CUDA_VISIBLE_DEVICES=0,1,2,3 vllm serve yandex/AliceAI-Foundation-80B-A3B-Base \
  --port 8001 \
  --trust-remote-code \
  --tensor-parallel-size 4 \
  --max-model-len auto \
  --attention-backend FLASH_ATTN \
  --attention-config.flash_attn_version=2

Docker-образ

Образ содержит сборку vLLM с поддержкой MTP для AliceAI. Требуются 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.

Для воспроизведения результатов бенчмарков удаляйте пробельные символы в конце полностью сформированного промпта с помощью prompt = prompt.rstrip() перед токенизацией.

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

Важно

Мы обнаружили ошибку, возникшую на этапе mid-training: значения смещений оценок экспертов (e_score_correction_bias) стали одинаковыми для всех экспертов. В результате эти смещения не влияют на выбор экспертов.

При дообучении этого чекпоинта установите коэффициент обновления e_score_correction_bias в 0, чтобы отключить их обновление.

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

Для подготовки агентских данных, использованных при обучении модели, мы использовали стандартный формат 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. Он загружает зафиксированную ревизию датасета tatsu-lab/alpaca, обучается только на ответах и сохраняет только LoRA-адаптер. Для модели такого размера необходим FSDP2, пример ниже рассчитан на четыре GPU с 80 GB памяти.

pip install \
  transformers==5.16.1 \
  accelerate==1.14.0 \
  peft==0.20.0 \
  datasets==5.0.1 \
  flash-linear-attention==0.5.0

pip install flash-attn==2.8.1 --no-build-isolation

CUDA_VISIBLE_DEVICES=0,1,2,3 \
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
accelerate launch \
  --use_fsdp \
  --num_processes 4 \
  --num_machines 1 \
  --dynamo_backend no \
  --mixed_precision no \
  --fsdp_version 2 \
  --fsdp_reshard_after_forward true \
  --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
  --fsdp_transformer_layer_cls_to_wrap AliceAIDecoderLayer \
  --fsdp_cpu_ram_efficient_loading true \
  --fsdp_sync_module_states true \
  --fsdp_state_dict_type SHARDED_STATE_DICT \
  finetune/finetune_lora.py \
  --model yandex/AliceAI-Foundation-80B-A3B-Base \
  --steps 100 \
  --sequence-length 512 \
  --output-dir alice-lora

--mixed_precision no здесь не означает FP32-модель: базовые веса загружаются в BF16, а LoRA-параметры PEFT хранит в FP32.

При RAM-efficient загрузке полные веса чекпоинта загружает только rank 0; остальные процессы создают модель на meta device и получают свои шарды через FSDP2.

---

Данная модель является предварительно обученной (pretrained) моделью и предоставляется в исходном виде, без дополнительного этапа post-training / alignment.

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

Перед использованием модели в production-среде рекомендуется провести собственное тестирование и, исходя из сценария применения, реализовать необходимые этапы дообучения, настройки и контроля поведения модели.

Пользователь самостоятельно определяет применимость модели для конкретного сценария и несет ответственность за ее интеграцию и использование.