Понимание параметров модели большого языка и требований к памяти: глубокое погружение

Вычисление количества параметров

Предварительно обученные модели для эффективного обучения трансформаторов

Расчет параметров в LLM на основе трансформатора

Давайте разберём расчёт параметров для каждого компонента модели LLM на основе Transformer. Будем использовать обозначения из оригинальной статьи, где d_model представляет собой размерность скрытых состояний модели.

1. Встраивание слоя:
   • Параметры = vocab_size * d_model
2. Многоголовое внимание:
   • Для пакетов h головы, с d_k = d_v = d_model / h:
   • Параметры = 4 * d_model^ 2 (для Q, K, V и прогнозов выхода)
3. Сеть прямой связи:
   • Параметры = 2 * d_model * d_ff + d_model + d_ff
   • где d_ff обычно 4 * d_model
4. Нормализация слоя:
   • Параметры = 2 * d_model (для масштаба и предвзятости)

Общие параметры для одного слоя Трансформера:

• Parameters_layer = Parameters_attention + Parameters_ffn + 2 * Parameters_layernorm

Для модели с N слои:

• Общие параметры = N * Parameters_layer + Parameters_embedding + Parameters_output

Пример расчета

Рассмотрим модель со следующими характеристиками:

• d_model = 768
• h (количество голов внимания) = 12
• N (количество слоев) = 12
• vocab_size = 50,000

1. Встраивание слоя:
   • 50,000 * 768 = 38,400,000
2. Многоголовое внимание:
   • 4 * 768^2 = 2,359,296
3. Сеть прямой связи:
   • 2 * 768 * (4 * 768) + 768 + (4 * 768) = 4,719,616 XNUMX XNUMX
4. Нормализация слоя:
   • 2 * 768 = 1,536

Общие параметры на слой:

• 2,359,296 4,719,616 2 + 1,536 7,081,984 XNUMX + (XNUMX * XNUMX XNUMX) = XNUMX XNUMX XNUMX

Общие параметры для 12 слоев:

• 12 * 7,081,984 = 84,983,808

Общие параметры модели:

• 84,983,808 + 38,400,000 = 123,383,808

Эта модель будет иметь около 123 миллионов параметров.

Типы использования памяти

При работе с LLM нам необходимо учитывать два основных типа использования памяти:

1. Модель памяти: Память, необходимая для хранения параметров модели.
2. Рабочей памяти: память, необходимая во время вывода или обучения для хранения промежуточных активаций, градиентов и состояний оптимизатора.

Расчет памяти модели

Память модели напрямую связана с количеством параметров. Каждый параметр обычно хранится в виде 32-битного числа с плавающей запятой, хотя в некоторых моделях используется обучение смешанной точности с 16-битными числами с плавающей запятой.

Память модели (байты) = количество параметров * байт на параметр

Для нашей примерной модели со 123 миллионами параметров:

• Память модели (32-битная) = 123,383,808 4 493,535,232 * 494 байта = XNUMX XNUMX XNUMX байта ≈ XNUMX МБ
• Память модели (16-битная) = 123,383,808 2 246,767,616 * 247 байта = XNUMX XNUMX XNUMX байта ≈ XNUMX МБ

Оценка рабочей памяти

Требования к рабочей памяти могут значительно различаться в зависимости от конкретной задачи, размера пакета и длины последовательности. Грубая оценка рабочей памяти во время вывода:

Рабочая память ≈ 2 * Память модели

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

Тренировочная память ≈ 4 * Модельная память

Для нашей примерной модели:

• Рабочая память вывода ≈ 2 * 494 МБ = 988 МБ ≈ 1 ГБ
• Тренировочная память ≈ 4 * 494 МБ = 1,976 МБ ≈ 2 ГБ

Устойчивое использование памяти и пиковое использование памяти

При обучении больших языковых моделей на основе архитектуры Transformer понимание использования памяти критически важно для эффективного распределения ресурсов. Давайте разделим требования к памяти на две основные категории: использование памяти в устойчивом состоянии и пиковое использование памяти.

Устойчивое использование памяти

Стационарное использование памяти включает в себя следующие компоненты:

1. Вес модели: копии параметров модели FP32, требующие 4N байт, где N — количество параметров.
2. Состояния оптимизатора: Для оптимизатора Адама это требует 8N байт (2 состояния на параметр).
3. Градиенты: копии градиентов FP32, требующие 4N байт.
4. Входные данные: Предполагая, что входные данные имеют тип int64, для этого требуется 8BD байт, где B — размер пакета, а D — входное измерение.

Общее использование памяти в устойчивом состоянии можно приблизительно определить следующим образом:

• M_steady = 16N + 8BD байт

Использование пиковой памяти

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

1. Нормализация слоя: требуется 4E байта на норму слоя, где E = BSH (B: размер пакета, S: длина последовательности, H: скрытый размер).
2. Блок внимания:
   • Вычисление QKV: 2E байта
   • Матрица внимания: 4 байта BSS (S: длина последовательности)
   • Выход внимания: 2E байта
3. Блок прямой связи:
   • Первый линейный уровень: 2E байта.
   • Активация GELU: 8E байт
   • Второй линейный уровень: 2E байта.
4. Кросс-энтропийная потеря:
   • Логиты: 6BSV байт (V: размер словаря)

Общий объем памяти активации можно оценить как:

• M_act = L * (14E + 4BSS) + 6BSV байт

Где L — количество слоев трансформатора.

Общее пиковое использование памяти

Пиковое использование памяти во время тренировки можно приблизительно определить, объединив устойчивую память и память активации:

• M_peak = M_steady + M_act + 4BSV байта

Дополнительный термин 4BSV учитывает дополнительное выделение в начале обратного прохода.

Понимая эти компоненты, мы можем оптимизировать использование памяти во время обучения и вывода, обеспечивая эффективное распределение ресурсов и повышение производительности больших языковых моделей.