ChatGPT, GPT-4 и Claude — это мощные языковые модели, которые дообучают, используя метод, который называется «обучение с подкреплением на основе отзывов людей» (Reinforcement Learning from Human Feedback, RLHF). Благодаря этому такие модели лучше отражают наши ожидания в плане их поведения, они лучше соответствуют тому, как мы собираемся их использовать.

В этом материале мы рассмотрим процесс обучения модели LLaMa c использованием RLHF. Модель будет учиться отвечать на вопросы с сайта Stack Exchange. В ходе обучения мы используем комбинацию из следующих методов:

• Дообучение с учителем (Supervised Fine-tuning, SFT).

• Моделирование вознаграждения / предпочтений (Reward / preference modeling, RM).

• Обучение с подкреплением на основе отзывов людей (Reinforcement Learning from Human Feedback, RLHF).

Комбинируя эти подходы, мы выпустили модель StackLLaMA. Эту модель можно найти здесь (вот — исходная модель LLaMA). Весь конвейер обучения модели входит в состав библиотеки Hugging Face TRL. Поэкспериментировать с моделью, понять — на что она способна, можно здесь.

Модель LLaMA

Применяя метод RLHF важно помнить о том, что в начале работы модель уже должна что-то уметь. Ведь на шаге применения RLHF занимаются лишь дообучением модели, её подстройкой под то, как с ней собираются взаимодействовать, и под то, как она должна отвечать на запросы пользователя. Поэтому мы решили воспользоваться недавно представленными эффективными моделями LLaMA. Это — самые свежие языковые модели, разработанные Meta AI. Их размеры варьируются от 7 до 65 миллиардов параметров, при их обучении используются наборы токенов в количестве от 1 триллиона до 1,4 триллионов. Это делает такие модели весьма функциональными. Мы, в качестве базовой модели, используемой на всех этапах нашей работы, применяем модель с 7 миллиардами параметров. Запросить доступ к исходному варианту этой модели можно здесь.

Набор данных Stack Exchange

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

Мы, назначая ответам баллы, следовали подходу, описанному в этой работе:

баллы = log2 (1 + лайки) округлённое до ближайшего целого числа, плюс 1 если автор вопроса принял ответ (мы назначаем −1 баллов в том случае, если количество лайков — это отрицательное число).

Позже мы увидим, что для модели наград всегда нужны два ответа на вопрос, которые можно сравнить. К некоторым вопросам имеются десятки ответов, что ведёт к множеству возможных пар ответов. Мы обращаемся, самое большее, к десяти парам ответов на один вопрос для того чтобы ограничить количество точек данных, связанных с одним вопросом. И наконец — мы приводим в порядок форматирование текста, конвертируя данные из формата HTML в формат Markdown, что позволяет улучшить читабельность выходных данных модели. Обсуждаемый тут набор данных и блокнот IPython, используемый для обработки данных, вы можете найти здесь.

Эффективные стратегии обучения моделей

Даже обучение самой маленькой модели семейства LLaMA требует огромного объёма памяти. Прикинем по-быстрому: при применении типа данных bf16 каждый параметр использует 2 байта (для fp32 это — 4 байта) в дополнение к 8 байтам, используемым, например, в оптимизаторе Adam (подробности смотрите в материалах по производительности трансформеров). Получается, что модели с 7 миллиардами параметров понадобится (2+8)*7B=70GB только чтобы поместиться в память. На самом же деле ей, вероятно, понадобится ещё больший объём памяти при расчёте промежуточных значений, таких, как показатели внутреннего внимания. Поэтому не получится, без дополнительных ухищрений, обучить подобную модель даже на чём-то вроде Nvidia A100 80GB. Можно, конечно, пойти на кое-какие хитрости, вроде использования более эффективного оптимизатора, или обучения с половинной точностью. Это позволит немного умерить аппетиты модели в плане памяти. Но, при решении подобных задач, памяти, рано или поздно, всё равно не хватит.

Ещё один подход к обучению таких моделей — это так называемое дообучение с эффективной обработкой параметров (Parameter-Efficient Fine-Tuning, PEFT). Соответствующие методики реализованы, например, в библиотеке peft, которая может проводить обучение низкоразмерных адаптеров (Low-Rank Adaptation, LoRA) на модели, загруженной в 8-битном формате.

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

При использовании такого сценария применяют следующее практическое правило: выделять примерно 1,2-1,4 Гб памяти на миллиард параметров (в зависимости от размера пакета и длины последовательности), чтобы в этот объём поместилось бы всё, что нужно для дообучения модели. Как сказано в материале, на который мы ссылались выше, это позволяет заниматься недорогим дообучением больших моделей (размерами до 50-60 миллиардов параметров на Nvidia A100 80GB).

Эти подходы позволяют осуществлять дообучение больших моделей на потребительских устройствах и в среде Google Colab. Среди примеров, достойных внимания, можно отметить дообучение facebook/opt-6.7b (13 Гб с типом данных float16) и openai/whisper-large в Google Colab (15 Гб памяти GPU). Для того чтобы узнать подробности об использовании библиотеки peft — обратитесь к нашему GitHub-репозиторию или к этому нашему материалу, посвящённому обучению модели с 20 миллиардами параметров на потребительском «железе».

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

Тут мы используем либо transformers.Trainer, либо accelerate. Обе эти библиотеки поддерживают распараллеливание данных без необходимости изменения исходного кода. Достигается это путём простой передачи аргументов при вызове скриптов с помощью torchrun или accelerate launch. Ниже показаны команды, где применяются эти два инструмента, с помощью которых обучающий скрипт запускают на 8 GPU, размещённых на одном компьютере.

accelerate launch --multi_gpu --num_machines 1 --num_processes 8 my_accelerate_script.py
torchrun --nnodes 1 --nproc_per_node 8 my_torch_script.py

Дообучение с учителем

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

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

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

Упаковка осуществляется с помощью ConstantLengthDataset. Затем, после загрузки модели с помощью peft, можно пользоваться Trainer. Сначала модель загружают с использованием типа данных int8, готовят её к обучению, а после этого добавляют адаптеры LoRA.

# загрузка модели с использованием 8-битного типа данных
model = AutoModelForCausalLM.from_pretrained(
        args.model_path,
        load_in_8bit=True,
        device_map={"": Accelerator().local_process_index}
    )
model = prepare_model_for_int8_training(model)

# добавление к модели адаптеров LoRA
lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)

model = get_peft_model(model, config)

Мы обучаем модель в течение нескольких тысяч шагов, используя цель CAUSAL_LM, и сохраняем модель. Так как мы снова будем дообучать модель, используя разные цели, мы объединяем веса адаптеров с весами исходной модели.

Заявление об отказе об ответственности: из-за особенностей лицензии LLaMA мы публикуем лишь веса адаптеров и контрольные данные модели. Вы можете подать заявление на доступ к весам базовой модели, заполнив эту форму, а после этого конвертировать их в формат трансформеров Hugging Face, воспользовавшись этим скриптом. Обратите внимание на то, что вам для этого, до выхода библиотеки версии 4.28, понадобится устанавливать соответствующие инструменты из исходного кода.

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

Моделирование вознаграждения и человеческие предпочтения

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

Тут мы прибегнем к одной хитрости, которая заключается в том, чтобы, вместо использования для обучения модели вознаграждений непосредственных отзывов, сделанных людьми, применить бы аннотации, сделанные людьми до цикла обучения с подкреплением. При таком подходе цель модели вознаграждений заключается в том, чтобы имитировать то, как люди оценивают тексты. Для построения модели вознаграждений можно прибегнуть к одной из нескольких стратегий. Самый простой способ — это предсказать отзыв человека (например — балльную оценку или бинарное значение, символизирующее «хорошо» или «плохо»). На практике же лучше себя показывает прогнозирование оценки для двух образцов, когда модель вознаграждений получает два образца-кандидата yk,yj для заданного ей промпта x, после чего ей нужно спрогнозировать то, какой из них будет выше оценён человеком.

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

Здесь — это оценка, назначенная моделью, а — это образец-кандидат, которому отдано предпочтение.

Работая с набором данных StackExchange, мы можем узнать о том, какой из двух ответов предпочёл пользователь, основываясь на оценке этих ответов. Имея эту информацию и функцию потерь, определённую выше, мы можем модифицировать transformers.Trainer, добавив в систему собственную функцию потерь.

class RewardTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        rewards_j = model(input_ids=inputs["input_ids_j"],  attention_mask=inputs["attention_mask_j"])[0]
        rewards_k = model(input_ids=inputs["input_ids_k"], attention_mask=inputs["attention_mask_k"])[0]
        loss = -nn.functional.logsigmoid(rewards_j - rewards_k).mean()
        if return_outputs:
            return loss, {"rewards_j": rewards_j, "rewards_k": rewards_k}
        return loss

Мы используем подмножество из 100000 пар кандидатов и оцениваем работу модели на наборе из 50000 пар образцов, которые не использовались при обучении модели. Используя скромный размер пакета учебных данных, равный 4, мы обучаем модель LLaMA с использованием LoRA-адаптера, созданного с помощью peft. Мы применяем оптимизатор Adam с точностью BF16. Вот наша конфигурация LoRA:

peft_config = LoraConfig(
    task_type=TaskType.SEQ_CLS,
    inference_mode=False,
    r=8,
    lora_alpha=32,
    lora_dropout=0.1,
)

Логирование при обучении осуществлялось с помощью Weights & Biases. Обучение длилось несколько часов, оно выполнялось средствами нескольких GPU Nvidia A100 на исследовательском кластере Hugging Face. Модель достигла итоговой точности в 67%. Хотя такой показатель и выглядит не очень высоким, надо отметить, что задача, которую мы решали, весьма сложна даже для человека.

В следующем разделе речь пойдёт о том, что адаптеры, после их обучения, могут быть объединены с «замороженной» моделью и сохранены для дальнейшего использования при решении похожих задач.

Обучение с подкреплением на основе отзывов людей

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

1. Генерирование ответов на основе промптов.

2. Оценка ответов с использованием модели вознаграждений.

3. Запуск шага оптимизации стратегии обучения с подкреплением с использованием оценок.

Промпты Query и Response приводят к следующему виду перед тем, как их токенизируют и передадут модели:

Question: <Query>

Answer: <Response>

Такой же шаблон оформления используется на этапах SFT, RM и RLHF.

Когда языковые модели обучают, применяя обучение с подкреплением, часто возникает проблема, которая заключается в том, что модель может научиться использовать модель вознаграждений в своих интересах. А именно, модель генерирует полную ерунду, что приводит к тому, что модель вознаграждений назначает тому, что получилось, высокую оценку. Для того чтобы это сбалансировать, мы пользуемся штрафом при расчёте вознаграждения. А именно, мы храним эталонную модель, которую мы не обучаем, и сравниваем то, что сгенерировала новая модель, с тем, что выдала эталонная модель, вычисляя расхождение Кульбака-Лейблера (Kullback–Leibler divergence, KL-divergence, KL):

Здесь — это вознаграждение, полученное от модели вознаграждений, а — это расхождение Кульбака-Лейблера между тем, что даёт текущая стратегия и тем, что даёт эталонная модель.

Повторимся — мы применяем peft ради эффективного использования памяти при обучении, что даёт нам дополнительные преимущества в контексте RLHF. Здесь эталонная модель и модель, где реализуется текущая стратегия, имеют одно и то же основание — SFT-модель, которую мы загружаем в 8-битном формате и «замораживаем» во время обучения. Мы оптимизируем, используя PPO, лишь LoRA-веса модели, реализующей текущую стратегию, используя при этом веса базовой модели.

for epoch, batch in tqdm(enumerate(ppo_trainer.dataloader)):
    question_tensors = batch["input_ids"]
        
    # Взятие образца из модели, реализующей текущую стратегию, и генерирование ответов
    response_tensors = ppo_trainer.generate(
        question_tensors,
        return_prompt=False,
        length_sampler=output_length_sampler,
        **generation_kwargs,
    )
    batch["response"] = tokenizer.batch_decode(response_tensors, skip_special_tokens=True)

    # Вычисление показателя тональности
    texts = [q + r for q, r in zip(batch["query"], batch["response"])]
    pipe_outputs = sentiment_pipe(texts, **sent_kwargs)
    rewards = [torch.tensor(output[0]["score"] - script_args.reward_baseline) for output in pipe_outputs]

    # Запуск шага PPO 
    stats = ppo_trainer.step(question_tensors, response_tensors, rewards)
    # Log stats to WandB
    ppo_trainer.log_stats(stats, batch, rewards)

Мы обучали модель в течение 20 часов на 3x8 GPU Nvidia A100 80GB, используя исследовательский кластер Hugging Face. Но достойный результат можно получить и гораздо быстрее (например — после примерно 20 часов на 8 GPU Nvidia A100). Все статистические данные о ходе обучения можно найти на Weights & Biases.

На что способна модель после обучения? Посмотрим!

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

Сложности, нестабильность, обходные пути

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

Правда ли то, что чем больше вознаграждение — тем выше эффективность?

Если в общих чертах рассматривать обучение с подкреплением, то можно сказать, что тот, кто им занимается, стремится достичь наивысшего вознаграждения. Мы, применяя RLHF, задействуем модель вознаграждений, не отличающуюся совершенством. Этим, в своих интересах, может воспользоваться алгоритм PPO. Проявиться это может в виде внезапного увеличения вознаграждения. Но когда мы изучали тексты, сгенерированные моделью, реализующей текущую стратегию, эти тексты, в основном, содержали повторы строки ```. Дело в том, что модель вознаграждений выяснила, что ответы на Stack Exchange, содержащие блоки кода, обычно получают более высокую оценку, чем ответы, в которых кода нет. К счастью, эту проблему мы замечали сравнительно редко, и, в целом, применение расхождение Кульбака-Лейблера должно способствовать противодействию подобным «выходкам» модели.

Правда ли то, что расхождение Кульбака-Лейблера — это всегда положительное значение?

Как мы уже говорили, показатель KL используется для того чтобы сдвинуть выходы модели ближе к выходам базовой модели. В целом, этот показатель измеряет расстояние между двумя распределениями, которое всегда выражается положительным числом. Но в trl мы используем оценочное значение этого показателя, которое, как ожидается, эквивалентно его реальному значению.

Очевидно то, что если токен взят из модели, реализующей текущую стратегию, для которой характерна более низкая вероятность, чем у SFT-модели, это приведёт к отрицательному штрафу KL, но, в среднем, он будет положительным, иначе нельзя было бы нормально получать образцы из модели, реализующей текущую стратегию. При этом некоторые стратегии генерирования данных могут привести к тому, что некоторые токены будут генерироваться, а некоторые будут подавляться. Например, при работе с пакетами завершённые последовательности подвергаются выравниванию, а при установке минимальной длины токены EOS подавляются. Модель может назначать очень высокие или очень низкие вероятности таким токенам, что приводит к отрицательным значениям KL. Так как алгоритм PPO осуществляет оптимизацию системы с учётом вознаграждения, он будет гоняться за этими отрицательными штрафами, что приведёт к нестабильности.

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

Существующие проблемы

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

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

Итоги

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

Воспользовавшись библиотекой peft любой может запустить наш пример на одиночном GPU! Если обучение будет слишком медленным, можно прибегнуть к распараллеливанию данных без изменения кода и масштабировать обучение путём использования нескольких GPU.

Если говорить о решении реальных задач, то это — лишь первый шаг! После того, как модель обучена, необходимо оценить её и сравнить её с другими моделями для понимания того, насколько хорошо она решает поставленную перед ней задачу. Сделать это можно, оценивая результаты генерирования данных разными версиями модели. Чем-то подобным мы занимались, создавая набор данных модели вознаграждений.

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

Мы активно работаем над улучшением TRL для того чтобы сделать все шаги, используемые при обучении модели с помощью RLHF, более доступными. Нам очень интересно узнать о том, что другие исследователи построят на основе наших разработок. Если вы хотите внести вклад в наше дело — загляните в раздел Issues этого GitHub-репозитория.