Салют, Хабр! В начале ноября мы делились с вами новостями о нашем флагмане GigaChat MAX и пообещали рассказать подробнее о процессе создания наших Pretrain-моделей. Пришло время сдержать слово и даже пойти дальше! 

Предобучение больших языковых моделей — это одна из наиболее ресурсозатратных стадий, которая непосредственно влияет на весь дальнейший процесс обучения GigaChat. От успешности обучения Pretrain-модели напрямую зависит качество всех следующих этапов обучения, например, Alignment и Vision. Поэтому сегодня мы хотим поделиться весами младшей модели линейки GigaChat версий base и instruct. Модель называется GigaChat-20B-A3B, так как построена на перспективной МоЕ-архитектуре!

Но и это ещё не всё. Вместе с весами мы делимся с сообществом улучшенной реализацией самого популярного класса МоЕ, а также кодом для механизма концентрации (а что это такое — читайте дальше ;)). Важно отметить, что хотя GigaChat-20B-A3B обучался на триллионах токенов преимущественно русского текста, он ещё способен на хорошем уровне понимать другие языки. Так что мы делимся мультиязычной моделью. О том, как запускать модель, какие версии доступны и как пользоваться контролируемой генерацией с помощью механизма концентрации, расскажем прямо сейчас!

Open source

Мы выкладываем на HuggingFace обе версии модели — base и instruct в fp32, обе с поддержкой длинного контекста в 131к токен. А также версии instruct модели в bf16 и квантизованную в int8.

P.S. Все замеры далее проводились в bf16

Воспользоваться моделями можно тремя способами:

• Через transformers. Взаимодействовать с GigaChat-ом можно будет с помощью кода ниже. Через cpu-offloading его можно запускать на любом устройстве

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig

model_name = "ai-sage/GigaChat-20B-A3B-instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, torch_dtype=torch.bfloat16, device_map="auto")
model.generation_config = GenerationConfig.from_pretrained(model_name)

messages = [
    {"role": "user", "content": "Докажи теорему о неподвижной точке"}
]
input_tensor = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt")
outputs = model.generate(input_tensor.to(model.device))

result = tokenizer.decode(outputs[0][input_tensor.shape[1]:], skip_special_tokens=False)
print(result)

• Через vllm. Для запуска быстрого и оптимизированного инференса GigaChat это наиболее простой и быстрый путь. У vllm есть ограничения на поддерживаемые GPU, мы замерялись на А100. Чтобы получить похожую скорость, можно арендовать ресурсы на cloud.ru.

from transformers import AutoTokenizer
from vllm import LLM, SamplingParams

model_name = "ai-sage/GigaChat-20B-A3B-instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
llm = LLM(model=model_name, trust_remote_code=True)
sampling_params = SamplingParams(temperature=0.3, max_tokens=8192)

messages_list = [
    [{"role": "user", "content": "Докажи теорему о неподвижной точке"}],
]

prompt_token_ids = [tokenizer.apply_chat_template(messages, add_generation_prompt=True) for messages in messages_list]

outputs = llm.generate(prompt_token_ids=prompt_token_ids, sampling_params=sampling_params)

generated_text = [output.outputs[0].text for output in outputs]
print(generated_text)

• Через vllm-сервер

# поднимаем сервер
vllm serve ai-sage/GigaChat-20B-A3B-instruct  \
  --disable-log-requests \
  --trust_remote_code \
  --dtype bfloat16 \
  --max-seq-len 8192

# делаем запрос из другого терминала
curl http://localhost:8000/v1/chat/completions -H "Content-Type: application/json"     -d "{
        \"model\": \"ai-sage/GigaChat-20B-A3B-instruct\",
        \"messages\": [
            {\"role\": \"system\", \"content\": \"Отвечай как ОЧЕНЬ безумный ученый\"},
            {\"role\": \"user\", \"content\": \"Что такое Mixture of Experts?\"}
        ]
    }"

Улучшенная реализация

К сожалению, в transformers не занесена модель с похожей на нашу архитектурой. Кстати, пару слов о ней будет прямо в следующем разделе! Поэтому мы были вынуждены использовать remote code.

В прошлой модели с похожей архитектурой выкладывались в репозитории deepseek-ai/deepseek-moe-16b-base. У этой реализации был набор проблем:

• если не установлен Flash Attention, то модель не запускалась и падала с ошибкой;

• неверный подсчёт RoPE;

• используется modelling от старой версии transformers (около 1 года назад).

Мы всё исправили и предлагаем обновленную реализацию. В нашем open source мы не стали переименовывать базовый класс, чтобы сохранить поддержку vllm и других open source-репозиториев.

Скорость генерации

Несмотря на большое количество параметров у модели, она получилась довольно быстрой. Мы провели тесты на одной видеокарте A100 с batch size = 1 в точности bf16. В качестве бэкенда использовался vllm. Наши замеры скорости можно воспроизвести через бенчмарк в репозитории vllm. Мы не тюнили параметры инференса вручную, а у MoE-слоёв есть такая возможность, поэтому остаётся запас по ускорению.

git clone https://github.com/vllm-project/vllm.git ./vllm_rep
cd ./vllm_rep/benchmarks
mkdir serving_bench

wget https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json

## выбираем модель для замера
# MODEL="meta-llama/Meta-Llama-3-8B-Instruct"
# MODEL="google/gemma-2-9b-it"
# MODEL="ai-sage/GigaChat-20B-A3B-instruct"
# MODEL="deepseek-ai/deepseek-moe-16b-chat"
# MODEL="Qwen/Qwen1.5-MoE-A2.7B-Chat"
# MODEL="Qwen/Qwen2.5-3B-Instruct"


## запускаем сервер и идем в другой терминал
vllm serve $MODEL \
	--disable_sliding_window \
	--disable-log-requests \
	--trust_remote_code \
	--dtype bfloat16 \
	--max-seq-len 8192

## в другом терминале
cd ./vllm_rep/benchmarks
### не забываем указать модель
SAMPLES_NUM=100
python benchmark_serving.py \
	--backend vllm \
	--model $MODEL \
	--dataset-name sharegpt \
	--dataset-path ShareGPT_V3_unfiltered_cleaned_split.json \
	--num-prompts ${SAMPLES_NUM} --seed 42 \
	--max-concurrency 1 \
	--ignore_eos \
	--save-result \
	--result-dir serving_bench \
	--result-filename $(basename $MODEL)_throughput_serving.json

Мы собрали в одном месте замеры скорости по популярным или архитектурно похожим моделям. Собранные в таблице метрики, получены скриптом из vllm:

• Total params: размер модели в миллиардах параметров;

• Active params: количество используемых миллиардов параметров на один forward;

• Req/s: обрабатываемое количество запросов в секунду;

• Output Token/s: генеративная пропускная способность токенов в секунду;

• Total Token/s: общая (input + output) пропускная способность токенов в секунду.

По скорости генерации новых токенов 20B MoE может немного уступать, но благодаря лучшей токенизации на русском языке (увы, замеры vllm проведены на английском) модель будет быстрее. Обратите внимание, что GigaChat 20В по скорости сопоставим с 3В-моделями, а по метрикам (про них будет ниже) — на уровне с 8В-моделями!

Кратко о том, что такое MoE

Сегодня тренд — обучать LLama-Like модели или, как говорят в статьях, трансформеры с пренормализацией и swiglu-функцией активации без использования dropout. Несмотря на это, мы решились на эксперимент и пошли немного иным путём, выбрав основой для наших моделей архитектуру «Mixture of Experts». Дальше будем её обозначать как MoE.

MoE отличается от LLama-Like архитектур изменённым MLP-блоком. Что это такое, мы подробно расскажем далее, а пока объясним это так. Классические Llama-like модели задействуют весь объём имеющихся у них знаний для каждой задачи. Но это не всегда эффективно. Ведь для решения задачи «Реши уравнение x^2 = 4. Чему равен x?» для модели нет необходимости вспоминать что-либо, кроме школьной математики. 

Эта идея и лежит в основе архитектуры. При получении определённого запроса модель избирательно активирует блоки с соответствующими знаниями (далее будем называть их экспертами). MoE-модель может иметь очень много экспертов, и её суммарный объём знаний оказывается довольно большим — это выливается в большой размер самой модели. Но благодаря выбору подмножества экспертов, стоимость использования сохраняется на уровне маленьких и менее вместительных по знаниям моделей.

Отлично звучит в теории! :) Но на практике возникают свои проблемы. Например, если в кругу экспертов появляются «рок-звезды», то они забирают всё внимание на себя. Другие эксперты в таком случае просто перестают использоваться и лежат в модели мертвым грузом. Этот момент нужно очень строго контролировать во время обучения. 

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

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

Переходим к deep learning

Важно отметить, что модель на самом деле выбирает экспертов не один раз, а много. Если у нас 32 слоя с 64 экспертами на каждом, а на forward выбирается по одному эксперту со слоя, то всего будет задействовано 32 эксперта из 2048 возможных (если бы мы на каждом слое выбирали всех экспертов)! И происходит это при КАЖДОЙ генерации нового токена. А при генерации десяти токенов будет сделано уже 320 выборов, то есть по 32 эксперта на каждый из десяти токенов. 

Мы много говорили про абстрактный выбор экспертов на слое, у нас этим занимается некий «роутер». Если коротко, то его задача — выбрать для каждого токена экспертов, которые его обработают. Если чуть сложнее, то в классическом варианте это обычный линейный слой (hidden_dim x num_experts) с softmax после него. После softmax мы берём индексы K наибольших логитов, то есть выбираем К экспертов на токен. В вышеприведённом примере на каждом слое выбираем одного эксперта из 64 возможных, и в примере мы берём того, у кого значение после роутера наибольшее. Собственно, вот и всё!

А как же выглядят сами эксперты? Мы уже намекали на это раньше: «MoE — это изменённый MLP-блок». То есть в классический MLP-блок на картинке ниже мы добавляем роутер и агрегацию экспертов. А вот каждый из экспертов представляется в виде обычных Gated MLP-блоков со swiglu-активациями.

Интерпретируемость экспертов

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

• Ноутбук с кластеризацией и эмбеддингами

• Ноутбук с управлением и генерациями

Метод концентрации

Как мы отметили выше, MoE-архитектура не использует все параметры модели при генерации ответов. Вместо этого она выбирает наиболее «подходящих» для входных данных экспертов. 

Поэтому возникает логичный вопрос: а за что именно эксперты отвечают? В ST-MoE были отдельные эксперты, ответственные за названия городов или пунктуацию, но в целом знания равномерно распределялись по сети. В Mixtral явной специализации экспертов и вовсе не было. Поэтому нам было интересно изучить поведение нашей модели, ведь у нас добавлены общие эксперты и скрупулёзно подобран гиперпараметр баланса (aux loss).

В своих экспериментах мы обнаружили, что эксперты нашей 20В-модели действительно специализируются на разных темах. Но самое интересное в том, что модель можно склонять к генерации ответов на определённые темы без всякого дообучения! Мы назвали это «методом концентрации». Мы как бы заставляем нашу модель «сконцентрироваться» на определённой теме.

Чтобы добиться такого эффекта, сначала необходимо понять, как выбираются эксперты, кто и за что отвечает. Для начала соберём статистику по выбору роутеров: будем использовать тексты на различные темы и фиксировать результаты их обработки. 

Мы знаем, какие эксперты активировались на каждом токене, в каждом слое, поэтому для всех экспертов модели и конкретной последовательности токенов мы можем составить эмбеддинг вида «9 2 4 5 6 4 8 2» из иллюстрации ниже. Эти цифры означают, что первый эксперт был выбран роутером 9 раз, второй эксперт — 2 раза и т. д.

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

Чтобы продемонстрировать процесс, воспользуемся набором Pile, который содержит данные из разных областей: программирования, медицины, веб-страниц и так далее. Из этого набора выберем 40 000 примеров и создадим для каждого из них эмбеддинг по методу выше. Затем эти эмбеддинги нужно будет кластеризовать, мы выбрали сочетание методов UMAP (2-d проекция) и HDBSCAN (кластеризатор). Возможны и другие подходы к кластеризации, но мы не углублялись в исследования по этому вопросу. Более подробные сведения вы найдёте в прикрёпленном ноутбуке.

Представляем вам результаты применения алгоритма UMAP. Как видно, примеры успешно распределились по кластерам согласно исходным темам набора Pile.

Поверх этой картинки применим кластеризатор и получим более мелкие кластеры. Осталось дело за малым: дать им интерпретацию.

Интерпретировать удобно с использованием векторов TF-IDF: они помогают выбрать десять ключевых слов, раскрывающих суть каждого кластера. Получившееся разделение оказалось весьма логичным: среди тем выделились программирование, алгебраические задачи, спорт, кулинария, география и другие. Например:

• кластер 7: teams, football, world, group, season, cup, team, week, games, round;

• кластер 25: like, food, coffee, just, cheese, wine, good, make, tea, chocolate;

• кластер 29: species, strains, isolates, genes, resistance, coli, cattle, gene, isolated, genus;

• кластер 51: room, property, home, bedroom, apartment, hotel, house, beach, area, located.

Рассмотрим, например, кластер 27: app, video, data, new, free, email, windows, use, software, using. Как вы помните, эмбеддинги получаются на основе выборов роутера. Тогда из них можно извлечь информацию о том, какие эксперты являются самыми важными для этого кластера, то есть содержат в себе информацию про интернет, сервисы, приложения и т. д. 

Эмбеддингом кластера назовём усреднение всех эмбеддингов текстов, входящих в этот кластер. С его помощью мы сможем выявить экспертов, которые чаще всего активируются на тексты этого кластера. Построим карту для этого эмбеддинга (heatmap): чем насыщеннее оттенок цвета, тем чаще использовался тот или иной эксперт. Можно заметить нескольких особенно выделяющихся экспертов — назовём их важными.

Карты активаций экспертов для текстов из кластера 27. По оси Y — слои модели, по оси X — номера экспертов. Чем ярче цвет, тем чаще выбирался эксперт. Слева представлены все активации, а справа — только для самых важных для этой темы экспертов.

По карте эмбеддинга справа можно построить его отфильтрованную версию, где ненулевые компоненты соответствуют важным для кластера экспертам. В нашем методе отфильтрованный эмбеддинг с некоторым весом прибавляется к логитам роутера до softmax, что позволяет увеличивать вероятности выбора конкретных экспертов. Это и есть концентрация: мы побуждаем MoE-модель использовать знания, то есть экспертов по определённой теме.

Применение

Для начала давайте попробуем сгенерировать что-нибудь без изменения поведения модели.

А теперь добавим управление на кластере 27: app, video, data, new, free, email, windows, use, software, using.

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

Попробуем теперь проделать то же самое, но возьмем кластер 25 про еду: like, food, coffee, just, cheese, wine, good, make, tea, chocolate

И все истории так или иначе стали про еду! Это хорошая демонстрация того, что MoE-эксперты обладают тематической специализацией.

Надеемся, что наше исследование помогло увеличить знания о том, что происходит внутри загадочных языковых моделей, а именно MoE. Рассматривая группы экспертов между разными слоями, мы смогли понять, за что они отвечают. А наш метод управления потенциально помогает добиться более управляемых генераций. Мы планируем провести больше экспериментов и рассказать о них в будущем. Также предлагайте свои модификации и тестируйте — you are welcome!

Метрики

А пока что повторим, что мы выкладываем две модели в открытый доступ — base и instruct. Мы проверили серию замеров: среди российского open source GigaChat занимает высокие позиции по многим метрикам.

Base-модель

Будет проще всего сравнить метрики среди моделей в размерах около 7 млрд параметров. Вероятно, можно было бы выделить больше моделей, но мы решили сравниться T-lite-instruct-0.1, LLaMa-3.1-8B и Gemma-2-9B, которые одновременно сильны по метрикам и способны неплохо говорить на русском. Qwen же не попал в этот короткий список, хотя его метрики с отрывом выше всех. Из проблем: что его base-, что instruct-версия норовит начать отвечать на китайском, независимо от языка запроса. Ещё один недостаток base-версии Qwen: малое количество русских токенов в токенизаторе, около пары тысяч.

В этом зачёте моделей GigaChat обходит LLaMa 3.1 и получает серебро. Не сомневаемся, что в задачах именно русскоязычного домена GigaChat на ваших дообучениях сможет превзойти Gemma-2-9B благодаря большому объёму русскоязычных тестов в обучающем наборе.

Что важно отметить: GigaChat в почти 2 раза дешевле по стоимости инференса, чем LLaMa 3.1 и обходит ее по качеству. Ускорение касается еще и обучения (более 60%, но подробнее об этом в следующем году), поэтому тюн Гигачата на любые задачи будет намного эффективнее и быстрее, нежели LLaMa based моделей. А еще по внутренним тестам производительности на 1 GPU наша архитектура держит больше потоков, то есть отлично подходит для высоконагруженных систем. Это для тех, кто захочет использовать GigaChat в промышленных решениях!

Instruct-модель

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

Во-первых, мы решили расширить табличку с метриками из статьи GigaChat MAX, теперь там есть все большие модели линейки GigaChat: Lite, Pro и Max.

Примечание по замеру метрик: GSM8k — это тест, который проверяет, как хорошо модели могут решать задачи с числами. В нашем исследовании мы использовали 5 шотов, чтобы оценить модель, и смотрели на последнее число в ответе. В оригинальное тесте ответ ищется по шаблону «### число». Тест Math тоже имеет разные версии, которые проверяют математические способности моделей. В нашем исследовании мы давали четыре примера и смотрели на последнее выражение в формате «\boxed{expression}». Затем оценивали результаты на совпадение с помощью библиотеки sympy.

Во-вторых, мы провели замер модели GigaChat-20B-A3B-instruct на бенчмарках MERA, Shlepa, ru-MMLU.

Внимание, наш open source «маленькой» модели догнал на MERA GigaChat Pro (получается, GigaChat ждёт новый апгрейд?), выбив средний балл 0,513. 

Итоги

Мы планируем регулярно делиться с вами результатами наших трудов. Уже сегодня открыли доступ к base- и instruct-версиям GigaChat, и надеемся, что они помогут вам достичь новых высот в решении ваших задач!

Мы стали первыми в России, кто опубликовал модель с MoE-архитектурой и поделился с вами интересными наблюдениями о ней. Особую гордость нашей команды вызывает то, что мы сумели соединить в одной модели очень высокую скорость инференса вместе с высоким качеством. А также предложили идею для исследования особенностей этой архитектуры методом концентрации!

Резюмируя:

• GigaChat-20B-A3B - первая открытая MoE модель в России. В мировой практике мало кому в мире удается обучать MoE архитектуры с хорошим качеством.

• GigaChat обучен преимущественно на русском языке, поэтому в отличие от тюнов Qwen и LLaMa моделей не совершает грамматических и пунктуационных ошибок и не переключается на другие языки во время разговора

• В GigaChat-20B-A3B дешевый инференс сочетается с хорошими показателями метрик

• Модель отлично подходит для исследований, например, концентрации

У нас складывается традиция — делать тизеры на будущее. Поэтому ожидайте продолжения истории о МоЕ, мы к вам ещё вернёмся с новым рассказом! Также хочется напомнить о нашем Telegram‑канале, где мы чаще делимся разными историями из мира ML и DL. 

Stay tuned! :) 

Благодарности

Хотелось бы сердечно поблагодарить:

• Команду pretrain за процесс исследований и оптимизаций новой архитектуры, обучение первой МоЕ-модели в России.

• Команду alignment за множество экспериментов для подготовки сильной instruct-модели.

• Команду RnD, без которой не было бы механизма концентрации.

• А также Валентина Мамедова (@vltnmmdv), Евгения Косарева (@evgenijkkk), Илью Щукина, Михаила Колесова, Дениса Мазура, Александра Прошунина и Артёма Червякова за помощь в подготовке open source и написании текста, а также за полезные замечания и исправления!