На 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": {...} // адаптеры

Я уже писал про библиотеку TSAD (Time Series Analysis for Simulation of Technological Processes) в канале, но отстраненно, как автор канала. Дам пару комментариев как соавтор и один из контрибьюторов проекта.

· Проект начинался как time series anomaly detection (поэтому TSAD), но потом разросся до time series analysis for diagnostics, и, наконец, решили не ограничиваться диагностикой, а пойти во все задачи промышленности (прогнозирование, оптимизация)
· Идея появилась, потому что просто не было open-source промышленных библиотек, а проблемы с предварительной обработкой данных, EDA, оценкой алгоритмов были
· Текущая цель библиотеки — сделать набор пайплайнов для AutoML в промышленности
· Функционала не достает, но есть несколько супер полезных функций/модулей, уже хорошо проработанных и реализованных:
1️⃣наборы данных для разных задач промышленности (RUL, CPD, outlier detection и др.)
2️⃣проработанная архитектура пайплайнов (самих пайплайнов бы побольше), что позволяет контрибьютить и разрабатывать сами пайплайны решения задач
3️⃣реализованные сложные метрики оценки типа NAB scoring algorithm, вот описание в документации. Их уже несколько лет активно используем в своих реальных проектах

Можно выдергивать и пользоваться этим функционалом в своих проектах.