Обучаем маленькую LLM с нуля на ноутбуке с 8 ГБ ОЗУ

Многие считают, что для обучения языковой модели с нуля нужна стойка с GPU. Это не так, по крайней мере если цель — разобраться. Бюджетного игрового ноутбука с Intel i5-9300H, 8 ГБ ОЗУ и GTX 1650 с 4 ГБ видеопамяти хватит, чтобы собрать маленький трансформер, скормить ему игрушечный датасет и наблюдать, как он становится лучше. Будет медленно и тесновато. Зато научит больше, чем любой облачный API.

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

Зачем это нужно

Есть три достойные причины взяться за это:

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

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

Следовать лекции, не копируя её

Естественный проводник здесь — серия видео Андрея Карпаты о создании GPT. Его лекция начинается с посимвольного токенизатора. Если вы уже понимаете токенизацию, можно забежать вперёд и сразу использовать byte-pair encoding (BPE). Смысл в обучении, поэтому разумно выбрать путь чуть сложнее показанного на экране и не перепечатывать его реализацию один в один.

Проект, а не ноутбук

Jupyter-ноутбуки удобны для экспериментов, но размывают структуру. Тому, кто живёт в Neovim, есть дополнительная причина от них отказаться. Вместо этого сделайте обычный Python-проект с обработкой ошибок, загрузчиком данных и каждым шагом в виде именованной функции. Каждая ячейка ноутбука становится одной функцией, а небольшой оркестратор вызывает их по порядку.

Порядок помогают поддерживать несколько решений:

1. Каталог .local хранит всё тяжёлое, что не стоит коммитить: сырые датасеты, токенизаторы и уже сконвертированные файлы токенов. Его задача — воспроизвести лучшее, что есть в сессии ноутбука, когда ничего не нужно загружать заново при каждом запуске.
2. Один конфиг-файл хранит настройки модели, пути и имена токенизаторов, так что в конвейере ничего не зашито жёстко.
3. Каждый шаг загрузки проверяет, существует ли файл, прежде чем скачивать его снова.

Данные и токенизатор

В качестве текста берётся датасет roneneldan/TinyStories с Hugging Face. Он небольшой, но не до смешного, а простая лексика подходит модели, которой приходится учиться на ограниченных вычислениях. Его предложил ассистент, так что выбор глубоко не исследовался, но для первого запуска он разумен. Функция-загрузчик скачивает его один раз и кеширует.

С токенизатором то же самое. Здесь это lakhera2023/gemma4-nano-tinystories, который загружается ещё одним маленьким условным загрузчиком. Опять предложение, принятое на веру и затем написанное руками, — и это честный способ учиться: пусть ассистент предлагает, а код вы пишете сами, чтобы понимать, что он делает.

Стена памяти

На токенизации машина с 8 ГБ начинает сопротивляться. Очевидный подход — токенизировать каждую историю, дописывать ID в один растущий массив в памяти и сохранить массив в конце. На ноутбуке с ограниченной ОЗУ такой подход намертво вешает всю систему.

Причина становится ясна, если присмотреться. Токен хранится как целое число, а список целых чисел занимает гораздо больше места, чем исходный текст. Накладные расходы списков в Python делают всё хуже, а многократное расширение большой структуры фрагментирует память. Через миллионы историй операционной системе уже ничего не остаётся.

Решение — вообще не держать результат в памяти. Записывайте токены каждой истории прямо в файл на диске, как только они получены:

with open(output_file, "wb") as f:
        for sample in dataset:
            text = sample["text"]
            tokens = tokenizer.encode(text)
            tokens.append(tokenizer.eos_id())
            np.asarray(tokens, dtype=np.uint16).tofile(f)

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

Сборка батчей

Батчинг проще, чем звучит, хотя требует небольшой разминки для ума. Языковая модель учится предсказывать следующий токен, поэтому целью для любого входа служит та же последовательность, сдвинутая на одну позицию вперёд. Возьмите срез токенов как x, срез, начинающийся на одну позицию позже, как y, и сложите много таких пар в батч. Для каждой позиции в x модель видит, что идёт следом, в y.

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

Что дальше

Часть с данными готова: закешированный датасет, закешированный токенизатор, файл токенов на диске и сборщик батчей. Следующий этап — сама модель, трансформер, написанный с нуля, всё так же по лекции Карпаты, с отступлениями там, где для них есть причина. Такой маленький ноутбук ограничит размер модели и длину контекста, так что настоящая работа будет в том, чтобы уместить что-то полезное в 4 ГБ видеопамяти.