group-telegram.com/senior_augur/348
Last Update:
На 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": {...} // адаптеры