AliceAI-Foundation-80B-A3B-Base
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.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
Также модель можно запустить через 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.
Для воспроизведения результатов бенчмарков удаляйте пробельные символы в конце полностью сформированного промпта с помощью prompt = prompt.rstrip() перед токенизацией.
Как дообучить под свои задачи
Формат данных
Для подготовки агентских данных, использованных при обучении модели, мы
использовали стандартный формат 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-среде рекомендуется провести собственное тестирование и, исходя из сценария применения, реализовать необходимые этапы дообучения, настройки и контроля поведения модели.
Пользователь самостоятельно определяет применимость модели для конкретного сценария и несет ответственность за ее интеграцию и использование.