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

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

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

Полный код тут. Если вкратце, то для запуска и экспериментов до 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":"Баланс аккаунта — это текущая сумма денежных средств, доступных в личном кабинете. Чтобы узнать точный баланс, я могу проверить его для вас."}]} …

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

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

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

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

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

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

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

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

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

Импорты и конфигурационные переменные с точки зрения реализации - тривиальны, поэтому пока на них не останавливаемся. Это просто импорты и набор констант для создания объекта 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 или правильности установки пакетов (это даст ошибку если, например, установлена CPU версия Torch).

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

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 здесь занимает 6.27 GB. Она неплохо говорит и помещается в домашнюю видеокарту, но gguf только для предсказаний - нет возможности подстройки.