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

Можно почитать документацию huggingface.

Виды операций на GPU

Tensor Contractions

Это самые ресурсоемкие операции в модели - например, линейные слои или механизмы внимания.

По сути все операции матричного умножения попадают в эту категорию.

GPU оптимизируют в основном эти виды операций.

Statistical Normalizations

Это операции, которые применяют некоторое отображение к входному тензору, например, softmax или layer normalization.

Такие операции менее требовательны к вычислительным ресурсам.

Element-wise Operators

Самые нетрудоемкие операции, применяются поэлементно, например, dropout, активации, residual связи.

Что влияет на память GPU

В процессе обучения модель использует гораздо больше памяти, чем при ее простом размещении на GPU, потому что во время обучения множество компонентов занимают память на GPU:

• Веса модели
• Градиенты
• Состояния оптимизатора
• Активации, сохраняемые для вычисления градиентов
• Временные буферы
• Память, специфичная для функциональности

Веса модели

• На один параметр модели с точностью fp32 приходится 4 байта;
• Для обучения в смешанной точности (mixed precision) приходится 6 байт, так как модель поддерживается в fp32 и в fp16, чтобы выполнять часть операций в полной точности, а часть в половинной;
• Для обучения полностью с точностью fp16, требуется 2 байта на параметр.

Хотя при смешанной точности (Mixed Precision) веса модели занимают больше памяти, чем при полной точности (FP32), оптимизация достигается за счет других факторов:

• скорость вычислений – операции в fp16 на современных GPU выполняются существенно быстрее
• экономия памяти на активациях и временных буферах - все это можем хранить только в fp16 по сравнению с весами
• снижение пропускной способности (bandwidth) памяти - фактически передаем меньше данных, поэтому вычисления на GPU ускоряются

Градиенты

Так как градиенты обычно всегда хранятся в fp32, даже для mixed precision (в mixed precision есть скейлер градиентов для конвертации их в fp32), то нужно 4 байта на параметр.

Состояния оптимизатора

• Для SGD без момента - ничего дополнительно хранить не надо
• Для SGD с моментом - храним 4 байта на параметр
• Для AdamW, например - храним 8 байтов на параметр, так как сохраняет не только накопленный средний градиент, но и квадрат градиента

В документации pytorch для SGD и AdamW можно увидеть, сколько и какие моменты используются во время градиентного шага.

Активации во время forward

Размер этих активаций зависит от длины последовательности, размера скрытого слоя, размера батча.

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

Временная память

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

Functionality-specific память

У программы могут быть особые требования к памяти, например, при генерации текста с помощью beam search, необходимо одновременно поддерживать сразу несколько вариантов генерации.