Оптимизация памяти для вывода и тонкой настройки больших языковых моделей

Большие языковые модели (LLM) như GPT-4, Bloom и LLaMA достигли замечательных возможностей за счет масштабирования до миллиардов параметров. Однако развертывание этих массивных моделей для вывода или тонкой настройки является сложной задачей из-за их огромных требований к памяти. В этом техническом блоге мы рассмотрим методы оценки и оптимизации потребления памяти во время вывода и тонкой настройки LLM на различных аппаратных конфигурациях.

Понимание требований к памяти

Память, необходимая для загрузки LLM, в основном определяется количеством параметров и числовой точностью, используемой для хранения параметров. Простое правило –

• Загрузка модели с X миллиардами параметров требует примерно 4X ГБ видеопамяти в 32-битной точности float
• Загрузка модели с X миллиардами параметров требует примерно 2X ГБ видеопамяти в 16-битной точности bfloat16/float16

Например, загрузка модели GPT-3 с 175 миллиардами параметров потребует примерно 350 ГБ видеопамяти в точности bfloat16. На данный момент самые крупные коммерчески доступные GPU, такие как NVIDIA A100 и H100, предлагают только 80 ГБ видеопамяти, что требует использования методов параллелизма тензоров и модели.

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

Количественная оценка памяти для вывода

Давайте оценим требования к памяти для вывода с помощью модели OctoCode, которая имеет около 15 миллиардов параметров в формате bfloat16 (~ 31 ГБ). Мы будем использовать Transformers library для загрузки модели и генерации текста:

from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch

<p>model = AutoModelForCausalLM.from_pretrained(&quot;bigcode/octocoder&quot;,
torch_dtype=torch.bfloat16,
device_map=&quot;auto&quot;,
pad_token_id=0)
tokenizer = AutoTokenizer.from_pretrained(&quot;bigcode/octocoder&quot;)
pipe = pipeline(&quot;text-generation&quot;, model=model, tokenizer=tokenizer)</p>

<p>prompt = &quot;Question: Please write a Python function to convert bytes to gigabytes.\n\nAnswer:&quot;
result = pipe(prompt, max_new_tokens=60)[0][&quot;generated_text&quot;][len(prompt):]</p>

<p>def bytes_to_gigabytes(bytes):
return bytes / 1024 / 1024 / 1024</p>

<p>bytes_to_gigabytes(torch.cuda.max_memory_allocated())

Вывод:

Пиковое использование памяти GPU составляет около 29 ГБ, что соответствует нашей оценке 31 ГБ для загрузки параметров модели в формате bfloat16.

Оптимизация памяти для вывода с помощью квантования

Хотя bfloat16 является распространенной точностью, используемой для обучения LLM, исследователи обнаружили, что квантование весов модели до более низкой точности, такой как 8-битные целые числа (int8) или 4-битные целые числа, может значительно снизить использование памяти с минимальной потерей точности для задач вывода, таких как генерация текста.

Давайте посмотрим на экономию памяти от 8-битного и 4-битного квантования модели OctoCode:

</div>
# 8-битное квантование
model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", load_in_8bit=True,
pad_token_id=0)
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
bytes_to_gigabytes(torch.cuda.max_memory_allocated())</pre>

Вывод:

# 4-битное квантование
model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", load_in_4bit=True,
low_cpu_mem_usage=True, pad_token_id=0)
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
bytes_to_gigabytes(torch.cuda.max_memory_allocated())

Вывод:

С 8-битным квантованием требование к памяти снижается с 31 ГБ до 15 ГБ, а 4-битное квантование снижает его еще больше до 9,5 ГБ! Это позволяет запускать модель OctoCode с 15 миллиардами параметров на потребительских GPU, таких как RTX 3090 (24 ГБ видеопамяти).

Однако обратите внимание, что более агрессивное квантование, такое как 4-битное, иногда может привести к ухудшению точности по сравнению с 8-битным или bfloat16 точностью. Существует компромисс между экономией памяти и точностью, который пользователи должны оценить для своего случая использования.

Квантование является мощным методом, который может позволить развертывание LLM в ресурсо-ограниченных средах, таких как облачные инстансы, устройства edge или даже мобильные телефоны, путем значительного снижения использования памяти.

Оценка памяти для тонкой настройки

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

Пиковое потребление памяти во время тонкой настройки обычно в 3-4 раза выше, чем при выводе, из-за дополнительных требований к памяти для:

• Градиентов
• Состояний оптимизатора
• Активаций из прямого прохода, сохраненных для обратного распространения

Консервативная оценка заключается в том, что тонкая настройка LLM с X миллиардами параметров требует примерно 4 * (2X) = 8X ГБ видеопамяти в точности bfloat16.

Например, тонкая настройка модели LLaMA с 7 миллиардами параметров потребует примерно 7 * 8 = 56 ГБ видеопамяти на GPU в точности bfloat16. Это превышает емкость текущих GPU, что требует распределенных методов тонкой настройки.

Распределенные методы тонкой настройки

Несколько распределенных методов тонкой настройки были предложены для преодоления ограничений GPU для крупных моделей:

1. Параллелизм данных: Классический подход параллелизма данных реплицирует всю модель на нескольких GPU, разделяя и распределяя пакеты обучающих данных. Это снижает время обучения линейно с количеством GPU, но не снижает пиковые требования к памяти на каждом GPU.
2. ZeRO Stage 3: Продвинутая форма параллелизма данных, которая разбивает параметры модели, градиенты и состояния оптимизатора на GPU. Это снижает использование памяти по сравнению с классическим параллелизмом данных, сохраняя только необходимые разбитые данные на каждом GPU во время разных фаз обучения.
3. Параллелизм тензоров: Вместо реплицирования модели параллелизм тензоров делит параметры модели на строки или столбцы и распределяет их на GPU. Каждый GPU работает с разбитым набором параметров, градиентов и состояний оптимизатора, что приводит к значительной экономии памяти.
4. Параллелизм конвейера: Этот метод разбивает слои модели на разных GPU/работниках, каждый из которых выполняет подмножество слоев. Активации передаются между работниками, снижая пиковое использование памяти, но увеличивая накладные расходы на связь.

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

Решение LLMem

Исследователи недавно предложили LLMem, решение, которое точно оценивает потребление памяти GPU при применении распределенных методов тонкой настройки к LLM на нескольких GPU.

Оценка использования памяти GPU для тонкой настройки предварительно обученных LLM

LLMem учитывает такие факторы, как рекомбинация параметров перед вычислением (ZeRO Stage 3), сбор вывода в обратном проходе (параллелизм тензоров) и различные стратегии распределения памяти для тела трансформера и головки языковой модели.

Экспериментальные результаты показывают, что LLMem может оценить пиковое использование памяти GPU для тонкой настройки LLM на одном GPU с ошибкой до 1,6%, превосходя среднюю ошибку DNNMem в 42,6%. При применении распределенных методов тонкой настройки к LLM с более чем миллиардом параметров на нескольких GPU LLMem достигает впечатляющей средней ошибки в 3,0%.