Estimating vRAM – Hamel’s Blog - Hamel Husain
Хэмел Хусейн делится практическими формулами для прикидочной оценки объёма vRAM, необходимого для обучения и инференса LLM. Он отталкивается от калькулятора Зака Мюллера и дополняет его поправками на количество обучаемых параметров, точность (dtype), оптимизатор Adam и накладные расходы прямого прохода. Ключевая формула для обучения с LoRA: к базовой оценке прибавляется (число обучаемых параметров в млрд × dtype/8 × 4), а итог умножается на 1.2. На примере: при 0.5B обучаемых fp16-параметров и базовых 14 ГБ выходит около 21.6 ГБ vRAM. Для инференса калькулятор подходит как есть, но нужно учитывать batch size при continuous batching. Также упоминаются способы экономии памяти: flash attention, gradient checkpointing и компиляция модели (например, MLC).
Оценки в этом посте — прикидки на салфетке. Они должны дать вам приблизительное представление о том, сколько vRAM вам нужно. Единственный способ узнать наверняка — запустить вашу модель на вашем железе и измерить использование памяти. Это особенно справедливо для обучения.1
Я буду обновлять этот пост по мере того, как буду узнавать больше и получать отзывы от других.
Предыстория
Мой друг Zach Mueller выпустил вот такой удобный калькулятор, который призван ответить на вопрос «Сколько памяти GPU мне нужно для модели?». Его калькулятор вбирает в себя всю математику, которая засела в головах у многих из нас, и упаковывает её в простой калькулятор.
Однако я уже пару недель обсуждаю с Zach и другими нюансы этого калькулятора и хочу поделиться дополнительной информацией, которая, как мне кажется, будет полезной.
Обучение с LoRA
Ключ к оценке памяти, необходимой для обучения, — отталкиваться от числа обучаемых параметров. Общая формула такова:
(Оценка из калькулятора в ГБ + ( число обучаемых параметров в миллиардах * (dtype обучаемых параметров / 8) * 4)) * 1.2
Вот обоснование для каждого из членов формулы:
dtype обучаемых параметров — это dtype ваших размороженных параметров модели для LoRA, обычно с иной точностью, чем остальная замороженная модель. Часто остальную часть модели загружают с 8-битной или 4-битной квантизацией, а адаптеры LoRA держат в 16- или 32-битной точности. Мы делим это на 8, чтобы получить количество байт на параметр. Например, если есть 1B обучаемых параметров и они 16-битные, это добавит 1 * (16 / 8) = 2GB к оценке, которую затем нужно умножить на другие величины (см. пример ниже). Множитель 4 привязан к популярному оптимизатору Adam (если вы используете другой, YMMV). Нужно 2x для оптимизатора, 1x для модели и 1x для градиентов. Множитель 1.2 — это дополнительный оверхед на прямой проход при обучении модели. Подробнее об этой эвристике можно прочитать в этом посте от Eleuther AI.
1, прибавленная к 20%, — это просто математический трюк для увеличения величины на процент. Например, если вы хотите увеличить значение на 20%, можно умножить на 1.20. Я упоминаю это только потому, что этот член формулы запутал некоторых читателей!
Пример
Например, если вы используете LoRA, 0.5B ваших параметров обучаемые и они в fp16, а калькулятор говорит, что вам нужно 14GB vRAM, то вот как вы рассчитаете общий объём памяти, необходимый для обучения:
(14GB + ( 0.5 * (16 / 8) * 4) ) * 1.2 =
(14GB + 4) * 1.2 = 21.6
Ответ: ~21.6 ГБ vRAM
Инференс
Калькулятор отлично подходит «как есть» для оценки vRAM, необходимой для инференса. Несмотря на то что есть другие оговорки, о которых стоит помнить, калькулятор — отличная отправная точка. Для инференса вам также придётся учитывать batch size для continuous batching, что зависит от конкретного сценария использования и от ваших требований к пропускной способности и задержке.
Оговорки
Есть и другие оптимизации, о которых стоит знать и которые могут уменьшить объём необходимой памяти:
Flash attention. Gradient checkpointing. Компиляция модели (например, MLC).
Распределённое обучение/инференс может добавлять дополнительный оверхед, но это сложная тема, которую я здесь не буду затрагивать.