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

Бенчмарки

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

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, 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. Он загружает зафиксированную ревизию датасета 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-среде рекомендуется провести собственное тестирование и, исходя из сценария применения, реализовать необходимые этапы дообучения, настройки и контроля поведения модели.

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