Улучшение производительности пропускной способности моделей Llama 2 с использованием Amazon SageMaker

Улучшение производительности моделей Llama 2 с Amazon SageMaker

Мы находимся в захватывающей точке в широком принятии машинного обучения (МО), и мы считаем, что большинство клиентских впечатлений и приложений будут переосмыслены с помощью генеративного искусственного интеллекта (ГИИ). ГИИ может создавать новый контент и идеи, включая разговоры, истории, изображения, видео и музыку. Как и большинство МО, ГИИ работает на основе моделей МО – очень больших моделей, обученных на огромных объемах данных и часто называемых базовыми моделями (БМ). БМ основаны на трансформаторах. Трансформаторы медленны и потребляют много памяти при генерации длинных текстовых последовательностей из-за огромного размера моделей. Большие языковые модели (ЯМ), используемые для генерации текстовых последовательностей, требуют огромных вычислительных мощностей и испытывают трудности в доступе к доступной памяти высокой пропускной способности (HBM) и вычислительной мощности. Это происходит потому, что большая часть доступной пропускной способности памяти потребляется при загрузке параметров модели и при авторегрессивном декодировании. В результате, даже при огромных вычислительных мощностях, ЯМ ограничены памятью ввода/вывода и ограничениями вычислений, не позволяющими полностью использовать доступные аппаратные ресурсы.

В целом, генеративное выводирование ЯМ имеет три основных проблемы (согласно Pope et al. 2022):

• Большой объем памяти из-за огромного количества параметров модели и временного состояния во время декодирования. Параметры часто превышают объем памяти одного акселераторного чипа. Кэши ключей-значений внимания также требуют значительного объема памяти.
• Низкая параллелизуемость увеличивает задержку, особенно с большим объемом памяти, требуя значительной передачи данных для загрузки параметров и кэшей в вычислительные ядра на каждом шаге. Это приводит к высоким требованиям к пропускной способности памяти для достижения целей задержки.
• Квадратичное масштабирование вычислений механизма внимания относительно длины последовательности усиливает задержку и вычислительные проблемы.

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

В этом посте рассматриваются техники для максимизации пропускной способности с использованием техник пакетирования для параллельного генеративного вывода в ЯМ. Мы обсудим различные методы пакетирования для уменьшения объема памяти, увеличения параллелизуемости и смягчения квадратичного масштабирования механизма внимания для увеличения пропускной способности. Цель состоит в полном использовании аппаратного обеспечения, такого как HBM и акселераторы, для преодоления узких мест в памяти, вводе/выводе и вычислениях. Затем мы подчеркиваем, как Amazon SageMaker large model inference (LMI) deep learning containers (DLCs) могут помочь с этими техниками. Наконец, мы представляем сравнительный анализ увеличения пропускной способности с каждой стратегией пакетирования на SageMaker с использованием LMI DLC для увеличения пропускной способности для моделей, таких как Llama v2. Вы можете найти сопровождающий пример в блокноте в репозитории примеров SageMaker на GitHub.

• Инновации для инклюзии Hack.The.Bias с помощью Amazon SageMaker
• Текст-в-музыку Генеративный ИИ Stability Audio, MusicLM от Google и многое другое
• Индустрия 4.0 Метавселенная разблокирована – как AR/VR, ИИ и 3D-технологии приводят следующую промышленную революцию.

Вывод для больших языковых моделей (ЯМ)

Авторегрессивное декодирование – это процесс, при котором языковые модели, такие как GPT, генерируют текстовый вывод по одному токену за раз. Он включает в себя рекурсивное подачу сгенерированных токенов обратно в модель в качестве части входной последовательности для предсказания последующих токенов. Шаги следующие:

1. Модель получает предыдущие токены последовательности в качестве входа. Для первого шага это начальный запрос, предоставленный пользователем.
2. Модель предсказывает распределение по словарю для следующего токена.
3. Выбирается токен с наибольшей предсказанной вероятностью и добавляется в выходную последовательность. Шаги 2 и 3 являются частью декодирования. На данный момент наиболее известные методы декодирования – это жадный поиск, поиск с лучом, контрастный поиск и выборка.
4. Этот новый токен добавляется во входную последовательность для следующего шага декодирования.
5. Модель повторяет эти шаги, генерируя один новый токен на каждом шаге, пока не будет произведен маркер конца последовательности или не будет достигнута желаемая длина вывода.

Сервирование моделей для ЯМ

Сервирование моделей для ЯМ относится к процессу получения входных запросов для генерации текста, делания выводов и возврата результатов в запрашивающие приложения. В этом процессе используются следующие ключевые концепции:

• Клиенты генерируют несколько запросов вывода, при каждом запросе состоящем из последовательности токенов или входных подсказок
• Запросы принимаются инференс-сервером (например, DJLServing, TorchServe, Triton или Hugging Face TGI)
• Инференс-сервер пакетирует запросы вывода и планирует пакет в выполнение, включающее библиотеки разделения модели (например, Transformers-NeuronX, DeepSpeed, Accelerate или FasterTransformer) для выполнения прохода вперед (предсказания последовательности выходных токенов) на генеративной языковой модели
• Выполнение генерирует выходные токены и отправляет ответ обратно на инференс-сервер
• Инференс-сервер отвечает клиентам сгенерированными результатами

Есть проблемы с планированием на уровне запроса, когда сервер вывода взаимодействует с исполняющим движком на уровне запроса, например, каждый запрос использует отдельный процесс Python, требующий отдельной копии модели, что ограничивает память. Например, как показано на следующей схеме, вы можете загрузить только одну копию модели размером 80 ГБ на экземпляр машинного обучения (ML) с общей памятью ускорителя 96 ГБ. Если вы хотите обслуживать дополнительные запросы одновременно, вам придется загрузить дополнительную копию всей модели. Это неэффективно по памяти и стоимости.

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

Техники пакетной обработки

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

Существует два основных типа пакетной обработки для запросов вывода:

• Клиентская (статическая) – Обычно, когда клиент отправляет запрос на сервер, сервер обрабатывает каждый запрос последовательно по умолчанию, что не является оптимальным для пропускной способности. Чтобы оптимизировать пропускную способность, клиент объединяет запросы вывода в одну пакетную нагрузку, а сервер реализует логику предварительной обработки для разделения пакета на несколько отдельных запросов и выполняет вывод для каждого запроса отдельно. В этом варианте клиенту необходимо изменить код для пакетной обработки, и решение тесно связано с размером пакета.
• Серверная (динамическая) – Другой метод пакетной обработки заключается в использовании вывода для достижения пакетной обработки на стороне сервера. Когда на сервер приходят независимые запросы вывода, сервер вывода может динамически группировать их в большие пакеты на стороне сервера. Сервер вывода может управлять пакетной обработкой для достижения указанного целевого времени отклика, максимизируя пропускную способность и при этом оставаясь в пределах желаемого диапазона времени отклика. Сервер вывода обрабатывает это автоматически, поэтому не требуется изменений кода на стороне клиента. Серверная пакетная обработка включает различные техники для дальнейшей оптимизации пропускной способности для генеративных языковых моделей на основе авторегрессивного декодирования. Эти техники пакетной обработки включают динамическую пакетную обработку, непрерывную пакетную обработку и пакетную обработку с пагинированием (vLLM).

Динамическая пакетная обработка

Динамическая пакетная обработка заключается в объединении входных запросов и отправке их вместе в виде пакета для вывода. Динамическая пакетная обработка является универсальной техникой пакетной обработки на стороне сервера, которая работает для всех задач, включая компьютерное зрение (CV), обработку естественного языка (NLP) и другие.

В контейнере LMI вы можете настроить пакетную обработку запросов на основе следующих параметров в serving.properties:

• batch_size – Определяет размер пакета
• max_batch_delay – Определяет максимальную задержку для агрегации пакета

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

Вы можете реализовать динамическую пакетную обработку на SageMaker, настроив serving.properties контейнера LMI следующим образом:

#Dynamic Batching
engine=Python
option.entryPoint=djl_python.huggingface
batch_size=64 #пример
max_batch_delay=1000 #пример
option.tensor_parallel_degree=2 #пример

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

Непрерывная пакетная обработка

Непрерывная пакетная обработка – это оптимизация, специфичная для генерации текста. Она улучшает пропускную способность и не жертвует латентностью времени первого байта. Непрерывная пакетная обработка (также известная как итеративная или скользящая пакетная обработка) решает проблему простоя GPU и основывается на подходе динамической пакетной обработки, дополнительно непрерывно добавляя новые запросы в пакет. На следующей диаграмме показана непрерывная пакетная обработка запросов. Когда обработка запросов 2 и 3 завершается, запланирован набор новых запросов.

Следующая интерактивная диаграмма подробно описывает, как работает непрерывная пакетная обработка.

(По материалам: https://github.com/InternLM/lmdeploy)

Вы можете использовать мощную технику, чтобы сделать модели языковых моделей (LLM) и генерацию текста эффективными: кэширование некоторых матриц внимания. Это означает, что первый проход по запросу отличается от последующих проходов. В первом проходе вам нужно вычислить всю матрицу внимания, тогда как для последующих проходов вам нужно вычислить только новое внимание к токену. Первый проход называется предварительной обработкой (prefill) во всей этой кодовой базе, а последующие проходы называются декодированием (decode). Поскольку предварительная обработка намного дороже декодирования, мы не хотим делать ее постоянно, но текущий выполняющийся запрос скорее всего выполняет декодирование. Если мы хотим использовать непрерывную пакетную обработку, как объяснялось ранее, нам нужно выполнить предварительную обработку в какой-то момент, чтобы создать необходимую матрицу внимания для объединения группы декодирования.

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

Вы можете настраивать следующие параметры в файле serving.properties контейнера LMI для использования непрерывной пакетной обработки:

• engine – Рабочий движок кода. Значения включают Python, DeepSpeed, FasterTransformer и MPI. Используйте MPI, чтобы включить непрерывную пакетную обработку.
• rolling_batch – Включает пакетную обработку на уровне итерации с использованием одной из поддерживаемых стратегий. Значения включают auto, scheduler и lmi-dist. Мы используем lmi-dist для включения непрерывной пакетной обработки для Llama 2.
• max_rolling_batch_size – Ограничивает количество одновременных запросов в непрерывном пакете. По умолчанию – 32.
• max_rolling_batch_prefill_tokens – Ограничивает количество токенов для кэширования. Это нужно настраивать в зависимости от размера пакета и длины входной последовательности, чтобы избежать нехватки памяти на GPU. Поддерживается только при использовании rolling_batch=lmi-dist. Наше рекомендация – установить значение на основе количества одновременных запросов x памяти, необходимой для хранения входных и выходных токенов для каждого запроса.

Ниже приведен пример кода для serving.properties для настройки непрерывной пакетной обработки:

# Непрерывная пакетная обработка
engine=MPI
option.entryPoint=djl_python.huggingface
option.rolling_batch=auto
option.max_rolling_batch_size=64 #пример
option.paged_attention=false
option.max_rolling_batch_prefill_tokens=16080 #пример
option.tensor_parallel_degree=2 #пример

Пакетная обработка PagedAttention

В процессе авторегрессивного декодирования все входные токены LLM создают свои тензоры ключа и значения внимания, и эти тензоры хранятся в памяти GPU для генерации следующих токенов. Эти кэшированные тензоры ключа и значения часто называются кэшем KV или кэшем внимания. Как указано в статье vLLM: Easy, Fast, and Cheap LLM Serving with PagedAttention, кэш KV занимает до 1,7 ГБ для одной последовательности в Llama 13B. Он также является динамическим. Его размер зависит от длины последовательности, которая является сильно переменной и непредсказуемой. В результате эффективное управление кэшем KV представляет собой значительную проблему. В статье установлено, что существующие системы тратят 60–80% памяти из-за фрагментации и избыточных резервирований.

PagedAttention – это новый алгоритм оптимизации, разработанный UC Berkeley, который улучшает процесс непрерывной пакетной обработки, позволяя кэш внимания (кэш KV) быть несвязанным, распределяя память в блоках фиксированного размера или страницах. Это вдохновлено концепциями виртуальной памяти и страничной организации, используемыми операционными системами.

Согласно статье vLLM, кэш внимания каждой последовательности токенов разбивается на блоки и отображается на физические блоки через таблицу блоков. Во время вычисления внимания ядро PagedAttention может использовать таблицу блоков для эффективного получения блоков из физической памяти. Это приводит к значительному сокращению потери памяти и позволяет использовать больший размер пакета, увеличенное использование графического процессора и более высокую производительность. На следующей рисунке показано разделение кэша внимания на несвязанные страницы.

На следующей диаграмме показан пример вывода с использованием PagedAttention. Основные этапы:

1. Получение запроса на вывод с входным промптом.
2. В фазе предзаполнения вычисляется внимание, и ключи-значения сохраняются в несвязанной физической памяти и отображаются на логические блоки ключей-значений. Это отображение сохраняется в таблице блоков.
3. Входной промпт проходит через модель (прямой проход), чтобы сгенерировать первый токен ответа. Во время генерации токена ответа используется кэш внимания из фазы предзаполнения.
4. При последующей генерации токенов, если текущий физический блок заполнен, выделяется дополнительная память несвязанным образом, что позволяет выполнять выделение памяти в нужный момент.

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

Следующий код является примером serving.properties для настройки пакетной обработки PagedAttention в контейнере LMI на SageMaker:

#Paged Attention Batching
engine=MPI
option.entryPoint=djl_python.huggingface
option.rolling_batch=auto
option.max_rolling_batch_size=64 #пример
option.paged_attention=true
option.max_rolling_batch_prefill_tokens=16080 #пример
option.tensor_parallel_degree=2 #пример

Когда использовать какую технику пакетной обработки

На следующей диаграмме представлены техники пакетной обработки на стороне сервера вместе с примером serving.properties в LMI на SageMaker.

В следующей таблице представлены различные техники пакетной обработки и их применение.

	Пакетная обработка PagedAttention	Непрерывная пакетная обработка	Динамическая пакетная обработка	Пакетная обработка на стороне клиента	Без пакетной обработки
Как это работает	Всегда объединять новые запросы на уровне токенов вместе с пагинированными блоками и выполнять пакетный вывод.	Всегда объединять новый запрос на уровне токенов и выполнять пакетный вывод.	Объединять новый запрос на уровне запроса; может быть задержка на несколько миллисекунд для формирования пакета.	Клиент несет ответственность за объединение нескольких запросов на вывод в одной пакетной нагрузке перед отправкой на сервер вывода.	При получении запроса выполнять вывод немедленно.
Когда это работает наилучшим образом	Это рекомендуемый подход для поддерживаемых моделей только декодера. Он подходит для рабочих нагрузок, оптимизированных по пропускной способности. Применимо только к моделям генерации текста.	Одновременные запросы, поступающие в разное время с одной и той же стратегией декодирования. Он подходит для рабочих нагрузок, оптимизированных по пропускной способности. Применимо только к моделям генерации текста.	Одновременные запросы, поступающие в разное время с одной и той же стратегией декодирования. Он подходит для рабочих нагрузок, требующих более высокую пропускную способность для времени отклика. Применимо к моделям CV, NLP и другим типам моделей.	Он подходит для случаев оффлайн-вывода, не имеющих ограничений по времени отклика для максимизации пропускной способности.	Редкие запросы на вывод или запросы на вывод с разными стратегиями декодирования. Он подходит для рабочих нагрузок с жесткими требованиями по времени отклика.

Сравнение пропускной способности разных методов пакетной обработки для большой генеративной модели на SageMaker

Мы провели сравнительное тестирование производительности модели Llama v2 7B на SageMaker с использованием контейнера LMI и разных методов пакетной обработки, описанных в этом посте, с одновременной обработкой 50 входящих запросов и общим количеством запросов 5 000.

Мы использовали три разных входных промпта разной длины для производительностного тестирования. В методах непрерывной и пакетной обработки PagedAttention длина выходных токенов для трех входных промптов была установлена на 64, 128 и 256 соответственно. Для динамической пакетной обработки мы использовали постоянную длину выходных токенов 128. Мы развернули конечные точки SageMaker для тестирования с типом инстанса ml.g5.24xlarge. В таблице приведены результаты тестирования производительности.

Модель	Стратегия пакетной обработки	Запросы в секунду на ml.g5.24xlarge
LLaMA2-7b	Динамическая пакетная обработка	3.24
LLaMA2-7b	Непрерывная пакетная обработка	6.92
LLaMA2-7b	Пакетная обработка PagedAttention	7.41

Мы видим, что использование пакетной обработки PagedAttention в сравнении с динамической пакетной обработкой для модели Llama2-7B на SageMaker с использованием контейнера LMI позволяет увеличить пропускную способность примерно в 2.3 раза.

Заключение

В этом посте мы объяснили разные методы пакетной обработки для вывода LLM-моделей и как они помогают увеличить пропускную способность. Мы показали, как методы оптимизации памяти могут повысить эффективность аппаратного обеспечения с помощью непрерывной и пакетной обработки PagedAttention и обеспечить более высокие значения пропускной способности, чем динамическая пакетная обработка. Мы видели увеличение пропускной способности примерно в 2.3 раза при использовании пакетной обработки PagedAttention в сравнении с динамической пакетной обработкой для модели Llama2-7B на SageMaker с использованием контейнера LMI. Вы можете найти записную книжку, использованную для тестирования разных методов пакетной обработки, на GitHub.