Настройка и обучение языковой модели (LLM) с использованием LoRA и PEFT

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

Минимальная конфигурация GPU, создаваемого Jupyter Server — 1 GPU Tesla A100 40Gb.

Шаг 1. Установка зависимостей

Перед началом работы необходимо установить необходимые библиотеки: pip, torch, transformers, accelerate, datasets, loralib, einops, scipy и sentencepiece. Они нужны для работы с нейронными сетями и обработки естественного языка.

!pip install -Uqqq pip
!pip install -qqq bitsandbytes torch transformerspeft peft \
    accelerate datasets loralib==0.1.1 einops==0.6.1 scipy sentencepiece

import os
import torch
import transformers
from datasets import load_dataset
from peft import (
    LoraConfig,
    PeftConfig,
    PeftModel,
    get_peft_model,
    prepare_model_for_kbit_training
)
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig
)


os.environ["CUDA_VISIBLE_DEVICES"] = "0"

Шаг 2. Загрузка базовой модели

Выберите подходящую предварительно обученную модель. В этом примере используется модель Intel/neural-chat-7b-v3-1, но можно выбрать другую, заменив значение переменной MODEL_NAME. Для сжатия модели, мы используем библиотеку BitsAndBytes:

MODEL_NAME = "Intel/neural-chat-7b-v3-1"

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

Создается объект конфигурации bnb_config, который определяет параметры для сжатия модели с помощью библиотеки bitsandbytes. Это позволяет эффективно использовать память и ускорить вычисления.

Используются следующие параметры:

load_in_4bit — загрузка модели в 4-битном формате с уменьшением ее размера в памяти;
bnb_4bit_use_double_quant — использование двойной квантизации для дополнительного уменьшения размера модели;
bnb_4bit_quant_type — тип квантизации ("nf4"), который определяет, как модель будет сжиматься;
bnb_4bit_compute_dtype — тип данных для вычислений (torch.bfloat16), который позволяет уменьшить потребление памяти и ускорить вычисления, сохраняя при этом точность.

model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    device_map="auto",
    trust_remote_code=True,
    quantization_config=bnb_config
)

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
tokenizer.pad_token = tokenizer.eos_token

Шаг 3. Добавление LoRA

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

def print_trainable_parameters(model):
  """
  Calculates and displays the total number of model parameters and the number of trainable parameters.
  """
  trainable_params = 0
  all_param = 0
  for _, param in model.named_parameters():
    all_param += param.numel()
    if param.requires_grad:
      trainable_params += param.numel()
  print(
      f"trainable params: {trainable_params} || all params: {all_param} || trainables%: {100 * trainable_params / all_param}"
  )

model.gradient_checkpointing_enable()

Этот метод активирует механизм контрольных точек градиента в модели. Благодаря этому экономится память при обучении, поскольку хранится только необходимая для вычисления градиентов информация.

model = prepare_model_for_kbit_training(model)

Функция prepare_model_for_kbit_training адаптирует модель для обучения с использованием определенного количества битов, что является частью оптимизации для уменьшения занимаемой памяти.

config = LoraConfig(
    r=8,
    lora_alpha=32,
    #target_modules=["query_key_value"],
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM"
)

Здесь создается конфигурация для LoRA (Low-Rank Adaptation), с помощью которой можно адаптировать только небольшую часть весов модели. Конфигурация имеет следующие параметры:

r и lora_alpha — параметры, контролирующие размер и мощность адаптации;
target_modules — список модулей модели, к которым будет применена адаптация LoRA;
lora_dropout — применение dropout к адаптированным весам;
bias — настройка использования смещения в адаптации;
task_type — тип задачи, для которой настраивается модель (в данном случае генерация текста).

model = get_peft_model(model, config)

Функция get_peft_model применяет PEFT (Parameter Efficient Fine-tuning) к модели, используя предварительно заданную конфигурацию LoRA. Это позволяет более эффективно провести тонкую настройку модели.

print_trainable_parameters(model)

Шаг 3. Дообучение модели

Прежде чем начать обучение, можно проверить, как модель генерирует ответы на основе заданного запроса.

Задание текстового запроса для модели

Ниже приведен многострочный фрагмент, задающий текстовый запрос для модели. Текст между <human> и <assistant> указывает на диалог между человеком и ассистентом. Метод .strip() удаляет начальные и конечные пробелы и переносы строк.

prompt = """
<human>: midjourney prompt for a girl sit on the mountain
<assistant>:
""".strip()

Параметры для генерации текста

max_new_tokens — максимальное количество новых токенов, которые модель может сгенерировать;
temperature — степень случайности в выборе слов, где меньшее значение приводит к более предсказуемому тексту;
top_p — вероятностный порог выбора слов; рассматриваться будут только слова с вероятностью выше этого порога;
num_return_sequences — количество возвращаемых последовательностей;
pad_token_id и eos_token_id — идентификаторы токенов заполнения и окончания предложения.

generation_config = model.generation_config
generation_config.max_new_tokens = 200
generation_config.temperature = 0.7
generation_config.top_p = 0.7
generation_config.num_return_sequences = 1
generation_config.pad_token_id = tokenizer.eos_token_id
generation_config.eos_token_id = tokenizer.eos_token_id

Преобразование запросов в формат, пригодный для модели

Запрос преобразуется в формат, пригодный для модели с помощью токенизатора. return_tensors="pt" указывает на то, что возвращаемые тензоры должны быть в формате PyTorch. Затем данные перемещаются на выбранное устройство (GPU).

device = "cuda:0"

encoding = tokenizer(prompt, return_tensors="pt").to(device)

Генерация текста без расчета градиента

В блоке with torch.inference_mode() происходит генерация текста без расчета градиентов, что уменьшает потребление памяти и ускоряет процесс. Метод model.generate() генерирует текст, используя заданные input_ids (токенизированный запрос), attention_mask (маска внимания для запроса) и ранее настроенные параметры генерации generation_config.

%%time
with torch.inference_mode():
  outputs = model.generate(
      input_ids = encoding.input_ids,
      attention_mask = encoding.attention_mask,
      generation_config = generation_config
  )

Декодирование ответа обратно в текст

Полученный ответ декодируется обратно в читаемый текст. skip_special_tokens=True указывает на то, что специальные токены, например, токен окончания предложения, должны быть пропущены при декодировании.

print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Шаг 4. Загрузка набора данных

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

Здесь используется функция load_dataset из библиотеки datasets для загрузки набора данных из CSV-файла. Файл midjourney_prompt_dataset.csv содержит данные, которые будут использоваться для обучения или тестирования модели. Функция load_dataset автоматически обрабатывает файл и конвертирует его в формат, удобный для работы с библиотекой.

data = load_dataset("csv", data_files="midjourney_prompt_dataset.csv")

Загрузка и токенизация запроса

Эта функция сначала генерирует текстовый запрос с помощью generate_prompt, а затем токенизирует его с помощью ранее загруженного токенизатора. Опции padding=True и truncation=True обеспечивают добавление паддинга (заполнение до единой длины) и обрезку длинных запросов соответственно.

def generate_prompt(data_point):
  return f"""
<human>: {data_point["User"]}
<assistant>: {data_point["Prompt"]}
""".strip()

def generate_and_tokenize_prompt(data_point):
  full_prompt = generate_prompt(data_point)
  tokenized_full_prompt = tokenizer(full_prompt, padding=True, truncation=True)
  return tokenized_full_prompt

Выбор подмножества данных для обучения

Здесь выбирается подмножество данных для обучения data["train"], после чего оно перемешивается методом shuffle(). Затем с помощью метода .map() к каждому элементу набора данных применяется функция generate_and_tokenize_prompt, которая преобразует и токенизирует данные. В результате получается подготовленный к обучению набор данных, где каждый элемент представляет собой токенизированный текстовый запрос.

data = data["train"].shuffle().map(generate_and_tokenize_prompt)

Шаг 5. Обучение

Настройте параметры обучения и начните процесс обучения модели.

Объект «TrainingArguments»

Здесь создается объект TrainingArguments, который содержит различные параметры для обучения:

per_device_train_batch_size — размер батча обучения на каждом устройстве;
gradient_accumulation_steps — количество шагов накопления градиента перед их обратным распространением;
num_train_epochs — количество эпох обучения;
learning_rate — скорость обучения;
fp16 — использование 16-битной точности с плавающей запятой для ускорения обучения и снижения потребления памяти;
save_total_limit — максимальное количество сохраняемых чекпоинтов;
logging_steps — частота логирования;
output_dir — директория для сохранения результатов обучения;
optim — оптимизатор; здесь используется 8-битная версия AdamW;
lr_scheduler_type — тип планировщика скорости обучения;
warmup_ratio — доля общего числа шагов обучения, в течение которых скорость обучения линейно увеличивается до заданной.

training_args = transformers.TrainingArguments(
      per_device_train_batch_size=1,
      gradient_accumulation_steps=8,
      num_train_epochs=4,
      learning_rate=2e-4,
      fp16=True,
      save_total_limit=3,
      logging_steps=1,
      output_dir="experiments",
      optim="paged_adamw_8bit",
      lr_scheduler_type="cosine",
      warmup_ratio=0.05,
)

Объект «trainer»

Trainer отвечает за процесс обучения модели и имеет следующие параметры:

model — модель, которая будет обучаться;
train_dataset — набор данных для обучения;
args — аргументы обучения, определенные выше;
data_collator — объект, который формирует батчи из данных.

Здесь используется DataCollatorForLanguageModeling, который подходит для задач языкового моделирования. mlm=False указывает, что обучение проводится без маскирования токенов (не MLM).

trainer = transformers.Trainer(
    model=model,
    train_dataset=data,
    args=training_args,
    data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False)
)

Отключение кэширования модели

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

model.config.use_cache = False

Метод «train()»

Метод train() объекта Trainer запускает процесс обучения модели с использованием заданных данных, аргументов и настроек.

trainer.train()

Шаг 6. Сохранение модели и инференс

После обучения сохраните модель и проверьте ее способность генерировать текст.

Метод «model.save_pretrained(«trained-model»)»

Сохраняет в папку trained-model текущее состояние модели, которое включает в себя веса модели и ее конфигурацию. Сохранение модели позволяет использовать ее позже без необходимости повторного обучения.

model.save_pretrained("trained-model")

Конфигурация PEFT

Здесь загружается конфигурация PEFT (Parameter-Efficient Fine-Tuning) из сохраненной модели. Далее создается новый экземпляр модели с помощью этой конфигурации. Параметры return_dict=True, quantization_config=bnb_config, device_map="auto" и trust_remote_code=True настраивают поведение модели, включая формат возвращаемых данных, настройки квантизации, автоматическое распределение по устройствам и доверие к исполняемому коду.

config = PeftConfig.from_pretrained('./trained-model')
model = AutoModelForCausalLM.from_pretrained(
    config.base_model_name_or_path,
    return_dict=True,
    quantization_config=bnb_config,
    device_map="auto",
    trust_remote_code=True
)

Загрузка токенизатора и установка токена

Загружается токенизатор, соответствующий модели, и устанавливается токен заполнения (`pad_token`) равным токену окончания предложения (`eos_token`), что часто необходимо для моделей генерации текста.

tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
tokenizer.pad_token = tokenizer.eos_token

Создание экземпляра модели PERT

Создается экземпляр модели PEFT из сохраненной модели, что позволяет использовать параметры, полученные в результате тонкой настройки.

model = PeftModel.from_pretrained(model, './trained-model')

Преобразование текстового запроса в формат тензоров PyTorch

Задается текстовый запрос, который затем токенизируется и преобразуется в формат тензоров PyTorch для отправки на вычислительное устройство.

generation_config = model.generation_config
generation_config.max_new_tokens = 200
generation_config.temperature = 0.7
generation_config.top_p = 0.7
generation_config.num_return_sequences = 1
generation_config.pad_token_id = tokenizer.eos_token_id
generation_config.eos_token_id = tokenizer.eos_token_id

device = "cuda:0"

prompt = """
<human>: midjourney prompt for a girl sit on the mountain
<assistant>:
""".strip()

encoding = tokenizer(prompt, return_tensors="pt").to(device)

Отключение вычисления градиентов для экономии ресурсов

В блоке with torch.no_grad() отключается вычисление градиентов для экономии ресурсов. Метод model.generate используется для генерации текста на основе входных данных (input_ids, attention_mask) и предварительно настроенных параметров генерации (`generation_config`).

%%time

with torch.no_grad():
  model.config.use_cache = False
  outputs = model.generate(
      input_ids = encoding.input_ids,
      attention_mask = encoding.attention_mask,
      generation_config = generation_config
  )

Декодирование и вывод первого элемента сгенерированного текста

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

print(tokenizer.decode(outputs[0]))