somedeadinsidespy/SleeplessLLaMa

Домашнее задание 3: Обучение LLaMA

Мягкий дедлайн: 11.11.2024 23:59
Жесткий дедлайн: 14.11.2024 23:59

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

О задании

LLaMA (Large Language Model Meta AI) – это большая языковая модель, занимающаяя на данный момент лидирующие позиции по качеству генерации. Как и у всех современных моделей, оригинальные слои Трансформера в ней заменены на улучшенные. Мы разбирали их всех на 5-ой лекции. Ваша задача будет состоять в том, чтобы самостоятельно реализовать архитектуру LLaMA, обучить ее на предложенных данных и замерить качество с помощью бенчмарка. При написании кода нельзя пользоваться сторонними ресурсами (даже с указанием источника). По итогам работы вам нужно будет написать отчет о том, что у вас получилось, а что нет, с какими сложностями вы столкнулись в процессе выполнения задания и как вы с ними справились.
Это задание сложнее всех предыдущих, поэтому на него дается больше времени. Мы настоятельно рекомендуем не откладывать его на потом. Мы не запрещаем оформлять код в ноутбуке, однако уверяем, что вам будет гораздо проще, если вы оформите задание в виде проекта.

Данные

Для обучения вам предлагается использоваться часть датасета OpenWebText. Это один из самых популярных датасетов для обучения трансформерных моделей. На нем обучались GPT-2, RoBERTa и многие другие. Полная его версия состоит из 8-ми миллионов текстов (38 GB), собранных по ссылкам из постов на Reddit. Все тексты уже почищены, поэтому делать с ними ничего не надо.

Так как у вас нет ни времени, ни ресурсов, чтобы учиться на всем корпусе, мы сократили его в 8 раз. Размер датасета рассчитан таким образом, чтобы вам не пришлось обучаться больше одной эпохи. Загрузить его можно следующим образом:

from datasets import load_dataset

dataset = load_dataset("ashaba1in/small_openwebtext")

Задания

Сформулируем все задания более подробно.

Внимательно прочитайте оригинальную статью LLaMA, особенно части про архитектуру и обучение.
Реализуйте модель по описанию со всеми архитектурными изменениями: RoPE, SwiGLU, RMSNorm, Pre-normalization. Опционально вы так же можете добавить ускорение attention c помощью xformers (гайд). Гиперпараметры вы можете подбирать по своему усмотрению, однако общий размер модели должен быть не меньше 100М параметров, не считая эмбеддингов. Ваша модель должна быть интегрирована в hugging face, чтобы ее можно было загружать и скачивать из Hub (об этом ниже).
Обучите модель на задачу безусловной генерации. В идеале ошибка на обучении должна опуститься ниже 4. Вы можете обучить свой токенизатор, но лучше взять готовый от MistralAI.

tokenizer = AutoTokenizer.from_pretrained('mistralai/Mistral-7B-v0.1')

Замерьте качество на бенчмарке tinyMMLU (об этом ниже).
Важно: в бенчмарке, который мы будем использовать, длина текстов лежит в диапазоне 512-1024 токенов. Поэтому вам надо сделать так, чтобы модель могла работать с такой длиной. У вас есть два варианта:
1. Обучить модель с длиной 1024 (Медленно, затратно по памяти)
2. Обучить модель с длиной 256 и дофайнтюнить с Position Interpolation. (Быстро, круто, чуть сложнее реализовать)

Тестирование

Для тестирования LLM обычно используются специально сконструированные бенчмарки. Бенчмарк – это сборник различных задач, проверяющих умение модели понимать текст и отвечать на вопросы. Наиболее популярные бенчмарки – MMLU, HellaSwag, WinoGrande, Arc и так далее. Посмотреть на полный список всех бенчмарков для разных задач можно тут, а почитать про то, чем они отличаются, тут.
Главная проблема всех этих бенчмарков в том, что они огромны. Всего лишь одно тестирование модели может стоить больше $10.000. Для действительно крупных моделей такие бенчмарки необходимы для точной и всесторонней оценки качества. Однако мы, очевидно, не можем позволить себе тестировать свои модели на них. Поэтому мы возьмем бенчмарк поменьше, а именно tinyMMLU. Он состоит всего из 100 вопросов, в которых модели предлагается выбрать один из 4 вариантов ответа. Для оценки вероятности ответа в модель подаются все 4 варианта по очереди с вопросом в виде префикса. По итогу оценивается метрика IRT (очень похожа на обычную точность). Подробнее можно прочитать в статье.

Как бенчить модель

Для запуска модели на бенчмарке мы будем использовать библиотеку lm-eval-harness. По ссылке можно найти все инстукции по установке. Для запуска теста удобнее всего использовать команду:

lm_eval --model hf --model_args pretrained=<hf_model_path> --tasks tinyMMLU --batch_size 4 --output_path <output_path> --log_samples --device cuda:0 --trust_remote_code

После ее вызова в <output_path> сохранится файл с выходами модели. Имея его, нужно запустить следующий код.

import json
import numpy as np
import tinyBenchmarks as tb


file_name = "<output_path>/<hf_model_name>/<file>.jsonl"

results = []
with open(file_path, 'r') as file:
    for i, line in enumerate(file):
        results.append(json.loads(line))

# Ensuring correct order of outputs  
results = sorted(results, key=lambda x: x['doc_id'])

y = np.array([float(item['acc_norm']) for item in results])

metric_results = tb.evaluate(y, 'mmlu')
print(metric_results)

Важно: Для запуска первой команды вам надо положить вашу модель на Hugging Face Hub. Для этого модель должна иметь определенный вид. Подробный гайд написан здесь. Обратите внимание, что все файлы вашей модели должны лежать в отдельной папке проекта. Без этого работать не будет. После загрузки модели на Hub, она должна подтягиваться с помощью следующего кода без ошибок.

model = AutoModelForCausalLM.from_pretrained(<model_name>, trust_remote_code=True)

trust_remote_code=True тут обязателен, так как код для модели лежит на хабе, а не в репозитории transformers. Если модель никак не хочет загружаться, то можно посмотреть на то, как выглядит формат GPT-2.

Мы не ждем от вас высоких значений метрик (все таки модель относительно маленькая и учится не очень долго). Хорошим достижением будет, если вы получите IRT > 0.3.

Подсказки

Обучать LLM не бывает легко, особенно в первый раз. Обучение занимает много времени (примерно, одну ночь), поэтому обратите внимание на следующие моменты.

Во время обучения логируйте все, что можете (ошибку, нормы градиентов, число токенов в батче, ...). Это поможет вам заметить, если что-то идет не так.
Не жгите ресурсы просто так. Перед запуском обучения тщательно проверьте, что все компоненты кода ведут себя так, как надо. Если модель не учится совсем, то у вас точно ошибка в коде, а не плохие гиперпараметры.
Не забывайте периодически сохранять модель и оптимизатор, чтобы можно было перезапустить обучение с последнего чекпойнта в случае чего.
Перед запуском долгого обучения проверьте, что все в порядке, обучив маленькую модель на маленькой выборке. Лосс должен падать, градиенты не взрываться и тд.
Не забывайте стандартные методы ускорения: accelerate, mixed-precision и тд. Подробный гайд тут.
При формировании батча регулируйте число токенов, а не число текстов. Лучше всего вытягивать все тексты батча в строку и нарезать на куски максимальной длины, чтобы не было паддингов. Тексты будут начинаться не сначала, но это повысит эффективность.
При тестировании гипотез старайтесь совершать одно изменение за раз. Постарайтесь, чтобы это изменение было обдуманным.

Оценивание

При оценивании мы будем смотреть на несколько факторов:

Ваш отчет. Опишите в нем подробно процесс выполнения задания. По отчету мы будем понимать, как вы думали и какие выводы делали из наблюдений.
Удалось ли обучить модель хоть как-то. Не забудьте добавить в отчет графики, в первую очередь лосс.
Организация кода. Код должен быть написан аккуратно, без дублирований и хардкода переменных.

Оценка = 0.5 * [правильность реализации модели и обучения] + 0.3 * [отчет] + 0.2 * [качество модели]

somedeadinsidespy
/

SleeplessLLaMa