SmolLM2: PPO

Для этой части взял SmolLM2-135M-Instruct и PPOTrainer из TRL 0.15.0. В reward и value загрузил обученную ранее reward-модель на основе SmolLM-135M-Instruct. Словари и ID токенов у двух баз совпадают, а шаблон чата для policy оставил от SmolLM2

Настройки обучения

Параметр Значение
GPU A100 80 ГБ
Train/test prompts 8163 / 2721
Эпохи по prompts 1
Rollout batch size 64
Rollout batches 128
PPO epochs 2
Learning rate 3e-6
KL coefficient 0,05
Максимальная длина генерации 256 токенов
Seed 42
Время ячейки обучения, включая сохранения 44 мин 26,637 с

При оптимизации использовал microbatch 8 и накопление градиента 4. Каждый rollout делился на 2 minibatches, для rollout forward batch был 16. BF16 включил, остальные настройки оставил такими: temperature 0,7, clipping policy/value 0,2, vf_coef 0,1, gamma 1, lambda 0,95

Всего получилось 128 rollout-батчей и 8192 эпизода. Это чуть больше 8163 train-промптов: TRL округляет число батчей вверх до 128 x 64. По схеме 2 PPO-эпохи x 2 minibatches на каждый rollout получается 512 обновлений оптимизатора

Результаты

Во время обучения reward-модель оценивала ответы, которые генерировала сама policy. Вот как менялись метрики; в последнем столбце среднее за шаги 119-128

Метрика Первый батч Последний батч Среднее последних 10
Оценка reward-модели 1,949578 5,470911 5,425047
KL к исходной policy 0,002095 22,881641 22,717114
Reward с учётом KL-штрафа 1,949473 4,326829 4,289192
Loss policy -0,000941 -0,006553 -0,003494
Loss value-модели 1,364650 15,944157 16,287682

Reward вырос, но вместе с ним выросли KL и loss value-модели. На последних 10 шагах policy/clipfrac_avg в среднем равен 0,008594, а val/clipfrac_avg - 1. По одному reward результат не оценить: PPO обучалась под эту же модель оценки. С числами из отдельного примера в README reward-модели их тоже не сравнивал, там другой формат чата и score берётся в другой позиции

Вероятности ответов

Проверил logprob на 256 chosen-ответах HumanLLMs с промптами из PPO train и на 256 ответах Dolly. Сами chosen-тексты PPO не училась повторять: во время обучения она генерировала свои ответы. Dolly для дообучения не использовал, его наличие в предобучении исходной модели не проверял

Logprob считался только по токенам ответа и EOS: промпт ограничен 256 токенами, ответ - 768. После PPO средний logprob на токен снизился на обоих наборах относительно исходной SmolLM2:

Ответы До PPO После PPO Средняя парная разница Медиана разницы
HumanLLMs chosen -1,761877 -2,102480 -0,340603 -0,328541
Dolly -2,439633 -2,668418 -0,228785 -0,142257
Ответы Средняя сумма logprob до После Средняя длина с EOS
HumanLLMs -434,845311 -505,622205 259,25
Dolly -208,627343 -219,698274 85,7578125

У 32 примеров Dolly обрезан промпт, у HumanLLMs - ни одного. В каждой группе обрезан один ответ до лимита 768 токенов с EOS

Logprob готовых ответов снизился на обоих наборах, хотя reward за время обучения вырос. Это разные проверки: в первом случае оценивается вероятность конкретного текста, во втором reward-модель оценивает сгенерированный ответ. Длины в таблице взяты из датасетов, по ним нельзя понять, какой длины ответы генерирует PPO

Пример ответа

Вопрос: What's your morning routine like?

Для обеих моделей использовались do_sample=False и max_new_tokens=256. Ниже убран только общий префикс system/user/assistant

До PPO:

My morning routine is quite simple, but it's essential for maintaining a healthy and productive work-life balance. Here's what I do every morning:

  1. Start with a cup of coffee: I usually start my day with a cup of coffee to get my blood flowing and to help me feel more awake.

  2. Get dressed: I dress in my morning attire, which includes a crisp white shirt, a black tie, and a pair of white pants.

  3. Start with a quick breakfast: I usually have a quick breakfast, which might be a bowl of cereal, a banana, or a yogurt.

  4. Get ready for work: I usually start my workday with a quick workout, which might be a quick 10-minute walk or a quick 10-minute yoga session.

  5. Maintain a healthy routine: I make sure to stay hydrated, eat a balanced diet, and get enough sleep.

  6. Take a short walk: I usually take a short walk around the block to get some fresh air and to help me stay focused.

  7. Get ready for the day: I usually start my day with a quick breakfast, a quick workout, and a quick breakfast.

После PPO: пустой ответ. После префикса assistant модель не вывела текст продолжения

До обучения модель хотя бы отвечала на вопрос, пусть и с повторами. После PPO здесь получился пустой ответ. На этом запросе стало хуже, несмотря на рост reward. Чтобы понять, как часто это происходит, нужно проверить больше вопросов

Downloads last month
252
Safetensors
Model size
0.1B params
Tensor type
F32
·
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Model tree for WhyZosa/llm-course-hw2-ppo

Finetuned
(390)
this model

Dataset used to train WhyZosa/llm-course-hw2-ppo