Эксперимент по подстройке Gemma 3 для вызова процедур

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

Далее термины “подстройка” или “тюнинг” взаимозаменяемы. Транслитерированные из английского термины плохо образуют формы слова. Возьмем задачу для примера. Пусть есть агент на базе Gemma-3-4b-it, и нам нужно сделать так, чтобы модель выдавала вызов процедуры, если во входном промте имеется смысл обращения к конфиденциальному функционалу агента, например – активен ли мой доступ, какие последние транзакции и т.п.

Полный код тут ^[1]. Если вкратце, то для запуска и экспериментов до 8-битной квантизации нужна машина не менее 24 ГБ GPU с виртуальным окружением (torch + requirements), CUDA 12.6-13.0. Качаем модель с HF, запускаем train_qlora.py и проверяем предсказания (текстовые) скриптом compare.py.

Успешно отработавшая подстройка модели должна через 10 минут (RTX 4090) выдать такой вывод (значение loss должно понижаться):

Далее – более детально.

Некоторые подготовительные сведения и датасет

Подстройка делается при помощи QLoRA. Для обучения тренируемых параметров нужен датасет. В нашем случае, датасет (examples.json) сгенерирован в Chat GPT и содержит 70 строк следующего вида:

…
{"messages": [{"role": "user", "content": "Мой доступ активен?"}, {"role": "assistant", "content": "<tool_call name="get_sensitive_data" arguments="{}"/>"}]}
…
{"messages":[{"role":"user","content":"Что такое баланс аккаунта?"},{"role":"assistant","content":"Баланс аккаунта — это текущая сумма денежных средств, доступных в личном кабинете. Чтобы узнать точный баланс, я могу проверить его для вас."}]}
…

Первая строка – пример целевого поведения ^[3], а вторая – пример того как модель ведет себя при запросах общего характера, в которые входят слова, но ответ не являются чувствительными данными. Этот формат называется Chat Messages Format (известный как формат сообщений OpenAI в рамках Chat Completions API ^[4]). Едва ли есть устойчивое название. Один из самых сложных моментов – необходимость препроцессинга: половину кода подстройки занимают формат-ориентированные преобразования. Нельзя вот прямо эти данные применять для обучения. JSON легко формировать и он годится как первоначальный источник, но далее придется работать с данными как-то преобразованными.

На самом деле, элементы словаря тренировочного датасета зависят от того, с какой целью делается подстройка модели. Для нашего случая, когда нужно чтобы модель запоминала диалоговую последовательность (фраза – инструмент), нужен такой формат где есть user и assistant, может быть другой формат json-файла примеров для доменной адаптации в специфической области знаний, для классификации, или для завершения фраз (entailment). Мы еще обсудим тренировочные данные позже при обсуждении маскирования и функции препроцессора.

Библиотеки, базовая модель

Для авторизации, загрузки модели, обучения, и инференса, нужен Torch. Ставится правильно подобранной под версию ОС и Cuda командой отсюда ^[5]. Нам нужны будут еще библиотеки (ставятся pip-ом по requirements.txt):

transformers, peft, bitsandbyte, datasets, accelerate, huggingface-hub

Я использовал локально скачанную вот эту модель Gemma 3 ^[6] (4 млрд. параметров). Для загрузки модели, надо уметь генерировать токен и предварительно авторизоваться на HF.

Оборудование (GPU)

Для обучения и инференса достаточно карты с 24 Гб видеопамяти и Cuda 12.6. Арендовал тут ^[7] 4090. Цена эксперимента в минимальном случае, с тремя подстройками (по 10 минут каждая) – две сотни рублей. Это на два часа времени. Из-за того что в образе виртуалки, который Вы предусмотрительно выберете, стоят дрова на видеокарту и стоит CUDA, время аренды сильно экономится. Главное, когда завершите, не забыть удалить и сервер и том дисковый. При недостаточных средствах, сервер сам отвалится, но не сразу, поэтому у меня однажды баланс в минус ушел на 2 т.р. Можно было этого не допускать. Карточки с большей памятью ^[8] на современной архитектуре довольно часто заняты, и, если сервер убрать на полку (shelve), потом сложно его запустить из-за недоступности карточек.

Скрипт подстройки модели

Импорты и конфигурационные переменные с точки зрения ^[9] реализации – тривиальны, поэтому пока на них не останавливаемся. Это просто импорты и набор констант для создания объекта peft.LoraConfig почти в самом конце программы. Для экономии памяти, размер пакета установлен равным одному примеру из датасета: BATCH_SIZE = 1. Размер пакета 2 тоже войдет в память (потребует 19,2 GB) это не изменит время, но на том же количестве эпох даст более высокие значения loss > 1.2.  Про настройки QLoRA (R, ALPHA, DROPOUT,TARGET_MODULES) будет информация в главе про QLoRA.

load_data(data_path):

Далее – load_data(data_path) – функция загрузки данных. Построчное чтение файла json.

preprocess_function(examples, tokenizer):

Следом определена функция препроцессинга, которая реализует маскирование фраз пользователя. Обучение – это процесс подбора весов с тем, чтобы по определенному входу был определенный вывод. Маскирование нужно для того, чтобы модель не училась выводить на всех токенах учебного примера, а только на тех что помечены "role": "assistant", что нужны. Мы не должны обучать модель предсказывать токены пользовательского ввода. Мы должны учить модель выдавать токены вызова процедуры в ответ на определенный пользовательский ввод. Поэтому нужно исключить из расчетов токены пользовательского ввода. Исключение пользовательского ввода происходит благодаря тому, что в PyTorch функция CrossEntropyLoss игнорирует метки -100, поэтому прямой ход обучения работает на всех данных (и пользовательский ввод и ответ LLM), а когда происходит подсчет невязки (loss), учитываются только те токены, у которых метка не -100.

В этой функции можно видеть что мы отдельно пробегаем по строкам примерам и делается форматирование (apply_chat_template() + tokenize=False) и затем вне цикла делается токенизация. Функция tokenizer( texts,... работает параллельно и выполняет паддинг не для каждого примера в отдельности. Для небольшого количества (70) примеров можно и не делить.

main():

Запуск подстройки модели происходит относительно просто. Вызываются определенные ранее функции для загрузки токенизатора, данных и их преобразования в нужный формат для библиотеки dataset и затем – пакетный (batched=True) препроцессинг элементов датасета.

Затем разделение данных на тренировочный набор и валидационный. Потом – проверка наличия GPU или правильности установки пакетов (это даст ошибку ^[10] если, например, установлена CPU версия Torch).

4-битная квантизация делается при помощи BitsAndBytesConfig. Вместо того, чтобы хранить каждый вес, как 4 байта (32 бита), мы храним веса в 4-битном виде. Алгоритм (тип) квантизации bnb_4bit_quant_type="nf4" может управлять тем как распределяется дискретные кванты по области определения распределения вероятности токенов. Выбранный тип NF4 имеет преимущества нормального распределения (тут теория ^[11] от авторов, тут практика от HF ^[12], тут на русском ^[13] языке). Или не ^[14] имеет. Лучше не стану много не писать про такие вещи. По роду деятельности мне сильно не хватает времени разбираться, я занят тем, что должен успеть внедрить ИИ (без разницы – куда и зачем). Но теория очень интересна. Математика ^[15] моих университетских изысканий в области численных методов и оптимального управления удивительно часто пересекается в общем представлении с машинным обучением.

bnb_4bit_use_double_quant=True еще больше экономит память, если кратко. Двойная квантизация описана в той же статье. Вообще вся квантизация для экономии памяти. При проблемах с корректностью или сходимостью здесь можно экспериментировать.

8-битная квантизация при обучении QLoRA адаптеров займет 17,8 GB вместо 12,8 GB и делается заменой соответствующих строк на эти:

…
# Configure 8-bit quantization
print("Configuring 8-bit quantization...")
bnb_config = BitsAndBytesConfig(
    load_in_8bit=True,
    llm_int8_threshold=6.0,
)

# Load model with quantization
print("Loading model with 8-bit quantization...")
model = AutoModelForCausalLM.from_pretrained(
    MODEL_PATH,
    quantization_config=bnb_config,
    device_map="auto",
    dtype=model_dtype
)
…

После настройки квантизации происходит загрузка модели (AutoModelForCausalLM.from_pretrained()). Параметр device_map позволяет управлять загрузкой слоев на GPU. У нас только одна карта, поэтому device_map="auto".

Далее, для считанного из хранящейся модели с диска токенайзера, устанавливается токен паддинга (заполнения) для последовательностей токенов разной длины. Этим токеном будет EOS (конец последовательности).

Далее отключается KV кеш модели (model.config.use_cache = False). Он не совместим с gradient_checkpointing=True. И теперь вещи, не особо относящиеся к подстройке модели закончились. Начинается неосредственно настройка параметров для QLoRA. Начнем новую главу.

QLoRA. Основная часть скрипта

Подстройка модели начинается здесь. prepare_model_for_kbit_training(model) выполняет подготовку модели. Чтобы показать количество тренируемых параметров используется print_trainable_parameters(). Покажет примерно вот такое:

Тренируется мизерное количество весов (0.7567%).

Далее создается массив параметров training_args = TrainingArguments(...). Для существенной экономии памяти применяется gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS.

Упрощенно матричное уравнение модели с примененным QLoRA адаптером на этапе подстройки выглядит так (в предсказании dropout не используется):

y = W·x + dropout((alpha/r) × B·A·x)

Здесь:

x – входной массив токенов,

y – выходной массив логитов для всех токенов известных модели (до преобразования в вероятности и семплирования),

W – оригинальная матрица весов модуля размер – d (зависит от модуля),

A – компонента разложенного массива тренируемых параметров размер [d(вход) × r]

B – компонента разложенного массива тренируемых параметров размер [r × d(выход)]

r– ранг (LORA_R = 16) характеризует емкость LoRA адаптера

alpha – масштабный множитель (LORA_ALPHA = 32), характеризует насколько сильно параметры адаптера влияют на предсказания модели.

dropout – функция обнуления случайно выбранных выходных значений LoRA адаптера. LORA_DROPOUT = 0.1 означает, что 10% будут обнулены.

TARGET_MODULES – список всех модулей модели, для которых будет действовать адаптер. В нашем случае, модулей довольно много. Бывает, что в подстройку включается только парочка. Каждый модуль имеет свою размерность. От этой размерности будет зависеть d в приведенной выше формуле. Если открыть configuration.json для базовой модели, то можно видеть размерность d в зависимости от модуля и компоненты (А или B) будет либо 2560 либо 10240:

…
  "text_config": {
    "hidden_size": 2560,
    "intermediate_size": 10240,
    "model_type": "gemma3_text",
    "num_hidden_layers": 34,
…

Результат подстройки 4-битного адаптера:

А в папке ./qlora_output появятся матрицы A, B (adapter_model.safetensors), масштабные множители и целевые модули (adapter_config.json).

Инференс. Использование модели

В качестве оценки модели, применяется скрипт compare.py с набором запросов (6 штук) 4 требуют вызова функции и еще 2 – запросы общего назначения. В них tool_call возникать не должен.

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

О параметрах сэмплинга. Если его нет (do_sample=False), то температура, top_k, top_p игнорируются. Поэтому этих параметров  нет в коде compare.py.

Наиболее правильно работает 8-битная квантизация при подстройке. В результате обучения, модель должна нормально отвечать на общий вопрос:

И должна вызывать функцию при запросе конфиденциальной информации:

Заключение

На карте RTX 4090 24 GB подстройка QLoRA модели на  32,788 млн. параметров из 4,3329 млрд. занимает 10 минут GPU и при квантизации 8 бит дает неплохие результаты (текстовая модальность) и помещается в 24 GB видео памяти с запасом (требует 17,930 GB). Предсказание требует больше. В карту с 16 GB не поместится.

4-битная квантизация (требует 12,858 GB) дает выигрыш в объёме памяти, но не двукратный, а модель проходит тест хуже и не годится для работы. По времени тренировки – те же 10 минут. Предсказание не квантованной моделью соответственно тоже в такой объем памяти не войдет.

Мерджинг весов и выдача точки доступа при помощи vLLM не рассмотрены, но если любопытно, в репозитории есть. vLLM закомментирована в requirements.txt.

Нужно отметить, что модель мультимодальная, и в статье не рассматривалось в полном объёме применение возможностей (изображения). Для предсказания с чистым текстом существуют модели куда менее требовательные к ресурсам. Например, квантованная на 6 бит gguf версия той же Gemma 3 здесь ^[16] занимает 6.27 GB. Она неплохо говорит и помещается в домашнюю видеокарту, но gguf только для предсказаний – нет возможности подстройки.