YandexGPT-5-Lite-8B-pretrain

YandexGPT-5-Lite-Pretrain

Pretrain-версия большой языковой модели YandexGPT 5 Lite на 8B параметров с длиной контекста 32k токенов. Обучение модели проходило в два этапа.

На первом этапе модель обучалась преимущественно на русскоязычных и англоязычных текстах общим объёмом 15T токенов с длиной контекста до 8k токенов. Состав датасета: 60% — веб-страницы, 15% — код, 10% — математика, остальное — другие специфичные данные, в том числе сгенерированная с помощью наших моделей синтетика и датасеты наших сервисов, например Яндекс Переводчика и база фактов Поиска.

На втором этапе, который мы назвали Powerup, модель обучалась на высококачественных данных объёмом 320B токенов. Состав Powerup-датасета: 25% — веб-страницы, 19% — математика, 18% — код, 18% — образовательные данные, остальное — синтетика, датасеты сервисов и прочие качественные тексты. На этом этапе мы увеличили длину контекста до 32k токенов.

Кроме того, наш токенизатор хорошо оптимизирован для русского языка. Например, 32k токенов нашей модели в среднем соответствует 48k токенам Qwen-2.5.

Более подробно — в нашей статье на Хабре.

Задавайте вопросы в discussions.

Бенчмарки

В своей категории модель достигает паритета с мировыми SOTA по ряду ключевых бенчмарков для pretrain-моделей, а по многим другим — превосходит их:

<img src="https://habrastorage.org/r/w1560/getpro/habr/upload_files/fab/0de/405/fab0de40517e1fd4efc1302eaaf325d8.png" alt="Таблица бенчмарков" width="100%"/>

* по данным репорта разработчиков модели. <br> BBH — 3-shot, HUMAN_EVAL и MPBB — 0-shot, все остальные бенчмарки — 5-shot. <br> Все замеры мы производили в HF transformers.

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

Модель можно запустить через HF Transformers:

from transformers import AutoModelForCausalLM, AutoTokenizer


MODEL_NAME = "yandex/YandexGPT-5-Lite-8B-pretrain"

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, legacy=False)
model = AutoModelForCausalLM.from_pretrained(
   MODEL_NAME,
   device_map="cuda",
   torch_dtype="auto",
)

input_text = "Кто сказал тебе, что нет на свете настоящей,"
input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")

outputs = model.generate(**input_ids, max_new_tokens=18)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Или через vLLM:

from vllm import LLM, SamplingParams


MODEL_NAME = "yandex/YandexGPT-5-Lite-8B-pretrain"

sampling_params = SamplingParams(
    temperature=0.3,
    max_tokens=18,
)

llm = LLM(
    MODEL_NAME,
    tensor_parallel_size=1,
)
input_texts = ["Кто сказал тебе, что нет на свете настоящей,"]
outputs = llm.generate(input_texts, use_tqdm=False, sampling_params=sampling_params)

for i in range(len(input_texts)):
    print(input_texts[i] + outputs[i].outputs[0].text)

Для полного соответствия токенизации мы рекомендуем пользоваться оригинальным sentencepiece:

import sentencepiece as spm
import torch
# git clone https://huggingface.co/yandex/YandexGPT-5-Lite-8B-pretrain
tokenizer = spm.SentencePieceProcessor(
    model_file="<path_to_local_repo>/tokenizer.model"
)
input_ids = tokenizer.encode(input_text, add_bos=True)
input_ids = torch.Tensor([input_ids]).to(model.device).to(torch.long)
outputs = model.generate(
    input_ids=input_ids,
    attention_mask=torch.ones_like(input_ids),
    max_new_tokens=18
)
print(tokenizer.decode(outputs[0].tolist()))

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

У нашей модели llama-like архитектура, это означает, что она совместима с большинством существующих фреймворков по дообучению LLM. Приведем короткий пример, как можно обучить нашу модель в torchtune:

Скачиваем репозиторий:

tune download yandex/YandexGPT-5-Lite-8B-pretrain \
  --output-dir YandexGPT-5-Lite-8B-pretrain

Смотрим список конфигов и копируем подходящий под задачу:

tune ls
tune cp llama3_1/8B_lora training_config.yaml

Изменяем конфиг, адаптируем его под нашу модель и задачу. Например, такой вариант подойдет для lora обучения на открытом инстракт датасете alpaca-cleaned.

Запускаем обучение:

tune run lora_finetune_single_device --config training_config.yaml

Подробности можно найти в официальной документации torchtune.