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

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

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

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

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

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

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

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

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

from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch
model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder",
torch_dtype=torch.bfloat16,
device_map="auto",
pad_token_id=0)
tokenizer = AutoTokenizer.from_pretrained("bigcode/octocoder")
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
prompt = "Question: Please write a Python function to convert bytes to gigabytes.\n\nAnswer:"
result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
def bytes_to_gigabytes(bytes):
return bytes / 1024 / 1024 / 1024
bytes_to_gigabytes(torch.cuda.max_memory_allocated())

Вывод:

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

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

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

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

</div>
# 8-bit quantization
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-bit quantization
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 с параметром 15B на потребительских графических процессорах, таких как RTX 3090 (24 ГБ видеопамяти).

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

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

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

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

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

• Градиенты
• Оптимизатор утверждает
• Активации из прямого прохода сохраняются для обратного распространения ошибки.

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

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

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

Для преодоления ограничений памяти графического процессора для больших моделей было предложено несколько распределенных методов точной настройки:

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

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

Решение LLMem

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

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

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

Экспериментальные результаты показывают, что LLMem может оценить пиковое использование памяти графического процессора для точной настройки LLM на одном графическом процессоре с частотой ошибок до 1.6%, превосходя среднюю частоту ошибок современного DNNMem, 42.6%. Применяя распределенные методы точной настройки к LLM с более чем миллиардом параметров на нескольких графических процессорах, LLMem достигает впечатляющего среднего уровня ошибок: 3.0%.