Как ускорить модели искусственного интеллекта, не перегружая память

Таблица ссылок

Аннотация и 1 Введение

2. Сопутствующая работа

3. Предыстория

3.1. Нотация и 3.2. Вывод языковой модели

3.3. Многозапросный, многоголовый и обобщенное многозапросное внимание

4. Контекстно-зависимое раздвоенное внимание и 4.1. Мотивация

4.2 Формулировка и 4.3 Сложность ввода-вывода памяти

5. Эксперименты

5.1 Сравнение возможностей Multi-Head, Multi-Query и Multi-Group Attention

5.2 Задержки моделей эквивалентных возможностей

5.3.Приложения

6. Заключение и ссылки

А. Часто задаваемые вопросы

Б. Сопутствующая работа

C. Установка

D. Многогрупповое внимание Семья

E. Контекстно-зависимое раздвоенное внимание

F. Приложения: Дополнительные результаты

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

Абстрактный

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

1. Введение

Появление больших языковых моделей (LLM) открыло новую эру машинного обучения, демонстрируя замечательную производительность при решении широкого спектра задач (Brown et al., 2020; OpenAI, 2023; Chowdhery et al., 2022; Touvron et al., 2023; Chen et al., 2021; Hoffmann et al., 2022; Li et al., 2022; Microsoft; Amazon, 2022; Nijkamp et al., 2023). Несмотря на их впечатляющие возможности, развертывание этих крупномасштабных моделей в практических приложениях создает значительные проблемы, особенно с точки зрения задержки вывода и эффективности. Улучшение этих аспектов имеет решающее значение, поскольку они напрямую влияют на вычислительные ресурсы, необходимые для генерации прогнозов, и обеспечивают практическую реализацию этих передовых моделей в различных отраслях.

Особенно требовательным сценарием вывода является выборка пакета с одним контекстом, где целью является генерация нескольких завершений из одного контекста. Эта задача обычно встречается в многочисленных приложениях, таких как инструменты IDE для редактирования кода, которые предоставляют несколько рекомендаций, или в случаях, когда для оптимальной производительности требуется ранжирование среди многих поколений (с помощью метрик ранжирования, таких как средняя логарифмическая вероятность, голосование большинства и т. д.). Инкрементальное декодирование такого сценария выборки требует интенсивного ввода-вывода памяти, что становится узким местом задержки для больших партий и длин контекста.

В этом исследовании мы изучаем две совместимые стратегии решения проблем ввода-вывода памяти в выводах преобразователей: (1) исследование многозапросности и ее компромиссов и (2) новый метод, называемый контекстно-зависимым раздвоенным вниманием.

Наше исследование начинается с анализа обобщенного внимания с несколькими запросами (Ainslie et al., 2023), которое включает в себя многозапрос (Shazeer, 2019), а также установленный механизм внимания с несколькими головами (Vaswani et al., 2017) для компромисса производительности и задержки. Наши результаты показывают плавное масштабирование производительности с увеличением размера модели для фиксированного значения числа групп g для обобщенного многозапроса[1]. Уменьшение g приводит к смещению вверх кривых масштабирования потерь при проверке в зависимости от размера модели. Последовательная связь между сжатием кэша, размером модели и потерями при проверке позволяет нам находить компромисс между эффективностью вывода и размером модели, т. е. позволяет нам выбирать более высокое сжатие для случаев использования, требующих высокой эффективности, при этом по-прежнему соответствуя производительности внимания с несколькими головами за счет компенсации большим размером модели.

Во-вторых, мы вводим контекстно-зависимое раздвоенное внимание, метод, который раздваивает любое внимание в обобщенном семействе многозапросов на контекстные и декодирующие компоненты во время инкрементального декодирования. Такое раздвоение включает в себя то же количество FLOP и дает идентичные результаты по сравнению с исходным вниманием, но может значительно снизить стоимость ввода-вывода памяти и, следовательно, задержку в сценариях с высокой длиной пакета и контекста. Этот подход позволяет генерировать несколько завершений в реальном времени без дополнительных затрат на задержку или позволяет использовать гораздо большие размеры пакетов, что приводит к улучшению производительности ранжирования. Например, для многоголовой модели CodeGen 16B (Nijkamp et al., 2022) с длиной контекста 2k мы можем увеличить размер пакета до 128 с раздвоенным вниманием по сравнению с размером пакета всего в 5 без него, в результате чего pass@k (Chen et al., 2021) увеличивается с 59,0% до 84,6%, а pass@top3 через среднее log-p увеличивается с 55,2% до 58,1%.

[1] Более низкие значения групп внимания g приводят к более высокому сжатию тензоров ключ-значение, как в случае с несколькими запросами, где g = 1, следовательно, повышая эффективность вывода и задержку из-за уменьшенного кэша KV по сравнению со случаем с несколькими заголовками, где g = h, количество заголовков внимания запроса.