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

🔹Числа с плавающей точкой (IEEE 754) — основной тип для вычислений в языковых моделях, у них есть знак, экспонента и мантисса. Экспонента контролирует диапазон значений, мантисса — точность в рамках этого диапазона. Также напомню, что есть приколы с представлением чисел около нуля (aka денормализованные числа). Есть куча реализаций разной битности:
— float: 32 бита, E8M23
— tf32: 19 бит, E8M10 (специфичный для Nvidia формат, отсюда все странности)
— fp16: 16 бит, E5M10
— bf16: 16 бит, E8M7 (экспонента и диапазон как у float)
— fp8: 8 бит, E4M3 или E5M2

🔹На что тратится память:
W: Сами веса модели
A: Активации (промежуточные состояния, результат вычисления функции от входа и весов)
G: Градиенты (обновления весов модели при обучении)
O: Состояние оптимизатора (моменты и дисперсия)

При инференсе есть только W и часть A, при обучении есть все 4 категории. Далее у каждого метода стоят буквы, которые обозначают, что именно экономится.

🔹Методы экономии памяти при инференсе:
— Квантование модели (WA): ужимаем тип данных для весов и активаций. В большинстве статьей так и написано: W4A16, что означает, что веса в 4 битах, активации в 16 битах.
— Flash Attention (A): оптимизируем вычисление внимания поблочными вычислениями в кэше GPU, попутно уменьшая сложность по памяти с квадратичной по длине последовательности до линейной.

🔹Дополнительные методы экономии памяти при обучении на одной карточке:
— Смешанная точность (A): имеем рабочую копию в 16 битах (bf16 или fp16), а также мастер-копию в 32 битах. Все операции делаем с рабочей копией и потом обновления весов вливаем в мастер-копию. Вы спросите: а где профит? А профит в том, что активации в 16 битах, а активации — это дофига памяти.
— Квантование оптимизатора (O): ужимаем тип данных для состояний оптимизатора. Чаще всего в 8 бит, перед собственно применением градиентов расквантовываем.
— Аккумуляция градиентов (AG): если мы хотим батч из больше чем одного примера, то A и G тоже раздуются. Но нам совсем не обязательно считать градиенты параллельно, поэтому мы можем считать их последовательно, суммировать, и только потом применять. Если это правильно😁 отмасштабировать, то это теоретически эквивалентно обучению на всём батче.
— Чекпоинты активаций (A): при обучении нам по-хорошему нужны активации на всех слоях, чтобы потом считать по ним градиенты. Но нам сначала нужно дойти до лосса, поэтому мы выкидываем часть промежуточных активаций и пересчитываем их заново на основе оставшихся чекпоинтов тогда, когда они нам реально понадобятся для подсчёта градиентов.
— Адаптеры (GO): основную модель вообще не трогаем, учим только новый маленький набор весов. Градиенты считаем только по нему, и на этом сильно экономим.

На практике используется буквально всё, везде и сразу 🤯
Типичный конфиг:

"model": {
    "attn_implementation": "flash_attention_2", // вы поняли
    "load_in_4bit": true, // квантование модели
    ...
},
"trainer": {
    "gradient_accumulation_steps": 32, // аккумуляция градиентов
    "bf16": true, // смешанная точность
    "optim": "adamw_8bit", // квантование оптимизатора
    "gradient_checkpointing": true, // чекпоинты активаций
    ...
},
"lora": {...} // адаптеры

The Pitfalls of Next-Token Prediction
Статья: https://arxiv.org/abs/2403.06963
Видео: https://www.youtube.com/watch?v=9V0bfZqT1Yo

Олды несомненно помнят, что в ранних seq2seq моделях, основанных на рекуррентных нейронных сетях, существовало два режима обучения: teacher-forcing, где на каждом шаге генерации в качестве входов использовались реальные токены, и другой режим с использованием токенов, предсказанных текущей версией модели. С появлением трансформеров и их параллельного обучения все стали использовать teacher-forcing. Авторы статьи возвращаются к этому вопросу.

🔹Задача
Авторы придумали простую синтетическую задачу: поиск пути между двумя вершинами в деревьях очень специфичной структуры, а именно в таких, где есть одна центральная вершина и несколько цепочек, исходящих из этой центральной вершины. Пример такого дерева (степень центральной вершины = 2, длина цепочек = 5):

8 ← 1 ← 5 ← 4 ← 3 → 0 → 2 → 6 → 7

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

Вход для задачи выглядит как случайно перемешанный набор рёбер дерева, плюс начало и конец пути (после "/"):

3 → 4 | 5 → 1 | 4 → 5 | 0 → 2 | 3 → 0 | 1 → 8 | 6 → 7 | 2 → 6 / 3 7 

Выход выглядит как сам путь:

3 → 0 → 2 → 6 → 7

Эту задачу мы решаем какой-нибудь моделью, которая умеет работать с последовательностями, например трансформером или рекуррентной сетью в авторегрессионном режиме (генерация токенов слева направо, как в языковых моделях).

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

Я потратил пару вечеров и воспроизвёл это в Колабе: ссылка. Воспроизводил для 2-5 деревьев, то есть ровно таких, как в примере выше. Код писал с нуля, но опираясь на их Гитхаб. Всё получилось, как написано в статье: усложнение задачи приводит к возможности её выучивания. Технически это выглядит просто как маскирование части input_ids.

🔹Про предсказание следующего токена
Щепотка "соломенного чучела": распространенная критика языковых моделей состоит в том, что они являются лишь "стохастическими попугаями", способными только предсказывать следующий токен. Считается, что из-за этого они не могут эффективно планировать или исправлять ошибки.

Однако авторы статьи предполагают, что основная проблема не в механизме предсказания следующего токена как таковом. Проблема — в teacher forcing'е, то есть в том, что во время обучения у модели нет необходимости планировать и пытаться сформулировать решение в активациях. И ведь большинство современных моделей обучалось именно с использованием этого метода.

🔹Ограничения
— Эмпирическая часть работает при фиксированном наборе гиперпараметров, и сломав их, можно сломать 2 и 3 наблюдение. Обучение и обучаемость таких моделей — это прежде всего оптимизационная задача, и там есть значительная доля случайности. Однако ни у меня, ни у авторов не получилось сделать модель, которая была бы контрпримером для первого наблюдения.
— У авторов нет никакого теоретического обоснования наблюдений. Как нет и алгоритма, по которому сеть считает путь. Мне кажется, что тут есть простор для творчества и механистической интерпретации.

У Anthropic пару недель назад вышел пост про агентов: https://www.anthropic.com/research/building-effective-agents

Он прекрасен тем, что определяет, что является агентом, а что не является. С точки зрения авторов поста, агент = система, в которой языковые модели динамически управляют собственными вызовами и инструментами, контролируя выполнение какой-то задачи.

Авторы утверждают, что для большинства случаев агенты не нужны: чем проще решение, тем лучше. С чем я полностью согласен 👏

Основное содержание поста — примитивы и паттерны оркестрирования языковых моделей без агентов. Основной примитив: улучшенная языковая модель, которая имеет доступ к инструментам, поиску и памяти. Этот примитив может быть реализован по-разному, например через конечное число последовательных вызовов языковой модели.

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

🔹Паттерн 2: маршрутизация
Если ваше приложение разбивается на несколько возможных параллельных путей, то стоит сделать классификатор, который будет определять нужный путь, и специализированные промпты под каждый из путей. Например, если вы делаете чатбот с несколькими независимыми функциями (рекомендация фильмов, ответы на вопросы по фильмам, чат на общие темы), то стоит использовать этот паттерн. В древних чатботах часто был детектор интентов, который делал ровно это 👴

🔹Паттерн 3: параллелизация
Если задача разбивается на несколько параллельных подзадач, то стоит их и вызывать параллельно. Например, если вам нужно извлечь огромный JSON из текста или переписки, возможно вам стоит извлекать его по кусочкам. Отличие от маршрутизации в том, что в ней нам нужна была только одна ветка, а тут нам нужны результаты всех вызовов.

🔹Паттерн 4: ведущий-ведомый 😭
То же самое, что и параллелизация, только с динамическим количеством и содержанием подзадач. Например, так можно делать агрегацию результатов поиска.

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

Ну и наконец последний паттерн — агенты, которые совершают действия в определенной среде, получают от среды обратную связь, и снова совершают действия.

Мне в разных местах в разное время пришлось использовать первые 3 паттерна. При этом тогда я не формулировал их как отдельные паттерны. Это не какие-то абстрактные штуки, это кристаллизация того, как удобно и просто строить системы (как и любые другие паттерны проектирования).

Это случилось быстро 😂

Статья с эпиграфом и буквицами 😎