Навчаємо маленьку 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 ГБ відеопам'яті.