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

Меня зовут Роман Горб, я старший ML-разработчик в команде YandexGPT. Тема инференса LLM заинтересовала меня, потому что я занимался R&D в квантовании сеток для CV-задач. Сегодня я расскажу, как безболезненно увеличить скорость инференса. Сперва разберёмся, зачем это нужно, а потом рассмотрим разные методы ускорения и фреймворки, которые могут в этом помочь. 

Эта статья основана на моей лекции на GPT Week, организованной Школой анализа данных. 

Что влияет на скорость инференса и зачем её увеличивать

Итак, на скорость инференса влияет несколько факторов. Самый очевидный — архитектура модели: Encoder Only, Decoder Only, Encoder-Decoder. Хороший пример разных типов моделей — GPT like-модель и T5 like-модель. Обе подходят для работы с текстами, но первая — это только декодировщик трансформера, а вторая — полный трансформер. Кроме того, размер модели и батча, с которым происходит инференс, тоже имеют значение. 

Второй фактор — тип задачи. Задачи могут быть генеративными или дискриминативными. К первым относятся Questioning and Answering и Summarization, а ко вторым — классификация и регрессия. Каждый генеративный сценарий, как правило, состоит из нескольких этапов: обработки контекста и самого декодинга, где мы токен за токеном генерируем ответ. На каждом из этапов ускорения можно добиваться по-разному. 

Дальше перейдём к нашему железному другу GPU. По сути, это асинхронный ускоритель с памятью и ядрами процессоров. Между ними есть шина, которая может перегонять данные. Если в какой-то момент мы упираемся в её пропускную способность и ядра оказываются недозагружены, наступает так называемое голодание (memory-bound). Из-за него падает производительность: операции выполняются с пропусками. Нужно стремиться избегать таких ситуаций и использовать ядра максимально эффективно. 

Взгляд со стороны пользователя: RPS и latency

Все факторы, которые я описал выше, всё же важнее для разработчиков. Для конечного пользователя значение имеют только две метрики: RPS и latency. RPS (requests per second) — это количество запросов, которые сервис способен обрабатывать в секунду, а latency — задержка ответа сервиса. Следовательно, для пользователя важно, чтобы показатель RPS был как можно больше, а latency — как можно меньше. 

Разберём на примере, как это устроено. Допустим, сервис должен обрабатывать 2000 запросов в секунду, а 20% из них мы можем просто закешировать, сэкономив ресурсы. Остаётся 1600 RPS, которые видеокарты должны обрабатывать в реальном времени. Отметим, что это не точная величина, ведь на пике активности значение RPS может подниматься и до 1800. При этом пользователь ожидает, что задержка при запросе составит не более трёх секунд — это целевая latency.

Тесты показали, что одна видеокарта способна держать 6 RPS, впритык укладываясь в заявленные три секунды. Получается, что для нашего сервиса нужно 300 видеокарт, но компания может предоставить только 200. Это ограничение, которое можно обойти, ускорив инференс. Таким образом, ускоряя инференс, мы можем добиваться лучших результатов меньшими ресурсами.

Однако, чтобы добиться желаемого, придётся пойти на компромисс между качеством модели и скоростью её работы. Выбрать подходящую модель поможет график ниже. На оси X находится время генерации в миллисекундах, но здесь можно отобразить RPS, время обработки контекста или генерации одного токена, количество токенов в секунду и другие важные метрики. На оси Y — относительное качество модели, то есть diff от какого-то абсолютного значения. 

Точки — это конкретные модели, из которых нужно выбирать. Откуда взять эти точки? Можно обучить несколько базовых моделей и нанести их на график, а дальше ускорять, добавляя всё новые и новые. Так на кривой появится огромное количество точек, а у нас — простор для выбора. 

Методы ускорения инференса 

Дистилляция

Обычно в этом методе есть какая-то большая модель, которую мы называем учитель (teacher), и модель поменьше — студент (student). Хорошим примером будет YandexGPT 3 — большая LLM, способная решать задачу с наилучшим качеством, но она совершенно не укладывается в наш вычислительный бюджет. Есть модель поменьше, вроде Т5, которая потребляет сильно меньше ресурсов, но не решает задачу так же качественно, как YandexGPT 3. Задача Knowledge Distillation состоит в том, чтобы минимизировать потери (loss) между фичами — предсказаниями учителя и студента. 

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

При дистилляции похожий процесс soft-label, но в objective добавляется ещё и вероятность, полученная от учителя. Под них будет мимикрировать студент. 

Давайте заменим функцию loss на KL-дивергенцию — показатель расстояния между двумя распределениями. Конечно, это не метрика в стандартном математическом понимании, но её всё равно можно использовать как функцию для минимизации. Вот как её рассчитать.  

Так как KL-дивергенция не считается аналитически, её нужно сэмплировать и оценивать по методу Монте-Карло. Чтобы не возникало неопределённости, знаменатель должен быть отличным от нуля, если числитель не равен нулю. Поэтому распределение студента должно накрывать распределение учителя. 

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

Решить эту проблему призвана обратная KL-дивергенция. Здесь мы сэмплируем уже по студенту в матожидании, меняя местами числитель и знаменатель. Такой подход закрывает возможность дифференцировать по сэмплированным значениям, но с этим можно бороться, причём разными методами. Один из методов включает использование MiniLLM.

Для дифференциации по обратной KL-дивергенции здесь ещё применяется и reinforce-трюк, подробнее о котором можно узнать в курсе Practical RL в ШАДе. Если коротко, то MiniLLM вводят аппроксимацию на градиенты, но те ведут себя нестабильно. Решить проблему можно с помощью PPO (Proximal policy optimization) — алгоритма, который использует методы оптимизации, чтобы обучать агентов. 

Этот метод подходит далеко не всем, ведь немногие специалисты умеют правильно настраивать алгоритм. Однако есть три трюка, которые способны заменить PPO.

Первый трюк. Нужно по-другому расписать градиент, чтобы получить формулу ниже. Также градиент можно расписать по шагам, вынести за скобку всё остальное и получить два слагаемых. Итоговая формула дифференцируема автодиффом. 

Второй трюк. Модели в современном NLP часто любят обманывать reward — взламывать его. Награды высокие, но сами модели не очень умные, поэтому тут нужно добавить регуляризацию, смешав текущие веса студента и изначальной модели. Чтобы просчитать градиенты, нужно поставить в формулу новую q, но это не так-то просто. Придётся прибегнуть к хаку Importance Sampling. 

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

В итоге всё выглядит следующим образом. Весьма похоже на дистилляцию soft-label.

MiniLLM превосходит все представленные бэйзлайны, в том числе KL и дистилляцию hard-label, а также обычное обучение без дистилляции. Если посмотреть на числа, то, например, при дистилляции 13B-модели в 1.3B-модель получится значение 60,7, а простое обучение 2.7B-модели даёт нам 55,4. Так что этот метод позволяет дистиллировать модель и получать качество лучше, чем даёт обычное обучение, даже для большего размера модели! 

Таблицы перечислены слева направо. 

• Таблица 1: Показатели ECE и точности на датасетах SST2 и BoolQ. Лучшие результаты студента выделены жирным шрифтом.

• Таблица 2: Отчётливые 4-граммы (Dist-4) и потери при моделировании языка (loss) на тестовых наборах на основе семейства LLaMA. MiniLLM сохраняет разнообразие.

• Таблица 3: Производительность на контрольном и тестовом наборах при применении различных комбинаций трюков в MiniLLM.

Все три трюка важны. Уберите один — и результат станет хуже. При этом трюки не вредят разнообразию — модель не беднеет в плане словарного запаса.

Квантизация

Следующий метод ускорения инференса — квантизация. Глобальная цель этого процесса — перевод весов из типов высокой точности вроде FP16 или FP32 в типы низкой, например, INT1/2/4/8 или, может быть, FP8. При такой операции уменьшается объём занимаемой памяти, но также растет и скорость инференса. Всё из-за стоимости низкобитных вычислений: если вы посмотрите спецификации современных видеокарт, то целочисленные перемножения матриц там стоят кратно дешевле. 

Обратите внимание на формулу ниже. X с плавающей точкой мы превратим в INT по хитрой формуле. Разделим его на некоторый scale, который подберём позже, возьмём от этого целую часть, добавим zero point, то есть какой-то сдвиг, чтобы распределение значений хорошо описывалось нашей формулой. После этого сделаем clip в какое-то из разрешённых значений — их всего два в степени количества битов. 

По такой формуле строится целочисленное представление вещественного числа, но её можно в каком-то смысле обратить — домножить на scale и вычесть zero point из результата. Так можно примерно восстановить исходное число, но, естественно, с потерей части информации. Эти операции называются квантизацией и деквантизацией соответственно, и с их помощью можно ускорять инференс.

Видов квантизации очень много. Во-первых, можно квантизовать только веса, а можно веса и активации, и получать разные типы ускорения. Наша формула — это один из самых простых методов, она использует равномерную сетку значений. Интовые числа расположены на прямой с одинаковым шагом, но можно располагать их и произвольно — такая квантизация называется Non-Uniform. 

Симметричная квантизация отличается  от асимметричной наличием zero point. Именно его добавление даёт сдвиг относительно нуля.

Если квантизовать каждый параметр с выделенным под него scale и zero point, никакого ускорения не получится. Поэтому нужно работать с целой группой значений — например, со строчкой матрицы или всем тензором. Вся группа при этом превращается в int-ы, сильно сжимается и обзаводится метаинформацией относительно небольшого размера. Такая схема уже способна давать прирост в скорости. 

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

Вообще, методы квантизации можно разделить на два больших класса. Первый — это Post-training quantization (PTQ). В этом процессе нет никакого обучения, а только подбор scale и zero point через статистическую калибровку на небольшой выборке. Для реализации этого метода достаточно будет 1000 примеров. Он стоит недорого и требует совсем немного GPU-времени. В противовес ему идёт Quantization-aware training (QAT). На него тратится колоссальное количество ресурсов, сотни и тысячи GPU-часов, но при этом качество получается выше.

Борьба за качество

Квантизация трансформеров началась с BERT like-моделей. Люди попробовали квантовать только веса, и всё получалось. А когда применили ещё и к активациям, возникли большие проблемы выбросов. В нашем тензоре появляются значения, которые сильно выпадают из распределения. При калибровке они учитываются в статистиках, из-за чего scale с zero point подбираются некорректно. Все методы, рассмотренные дальше, с этим борются. 

LLM.int8()

Первая статья, в которой проведено хорошее исследование по теме выбросов — LLM.int8(). Авторы в ней выяснили, что выбросы возникают в некоторых фиксированных слоях трансформера. Причём эмпирически доказано, что они не возникают при увеличении размера модели. На самом деле на это влияет падающая перплексия — метрика понимания текста. 

В таблице видно, насколько сильно меняются метрики, если убрать выбросы. 

SmoothQuant

SmoothQuant принадлежит к подклассу, который квантизует активации и веса. Как мы уже поняли, квантизовать активации довольно сложно из-за выбросов. Разумно частично переложить эту заботу в веса, тогда квантизация активаций станет чуть проще, а весов — чуть сложнее, но благодаря этому балансу в итоге у нас получится квантизация без потерь в качестве.

Рассмотрим, как выглядят типичные тензоры в настоящей LLM. Слева изображён тензор активаций, в котором видны большие пилообразные формирования — те самые выбросы. Они возникают несистемно в каналах тензора активаций, поэтому их невозможно определить раз и навсегда. Чтобы прийти к ситуации, изображённой слева, нужно сглаживание (smoothing). Как же его провернуть?

Нужно во все матричные перемножения в сети — синие блоки на схеме, кроме ReLU, — добавить переливание сложности. Оно представляет собой простое домножение на диагональную матрицу активаций и умножение на обратную к ней матрицу весов. Таким образом, при перемножении в линейном слое эти диагональные матрицы сокращаются. Определить диагональную матрицу можно с помощью специальной формулы, основанной на максимальных значениях матрицы активаций и весов смешанных по формуле со слайда. Альфа является гиперпараметром и просто подбирается перебором.

Просадки по качеству будут, но небольшие. В таблицах ниже проведены замеры бенчмарков для достаточно архаичных моделей. Наши замеры показывают, что качество сохраняется и на свежих моделях вроде YandexGPT. 

Существует несколько модификаций SmoothQuant, которые проиндексированы версиями О1, О2 и О3. В первой мы делаем квантизацию весов per-tensor (с одними scale на весь тензор) и per-token квантованием в активациях. Метаинформации при таком подходе больше, поэтому квантизовать проще. 

Мы в Яндексе провели замеры и выяснили, что, несмотря на меньшую скорость, О1-версия — самая качественная и единственная способна квантизовать без потерь. В остальных версиях потери существенные. Позже это подтвердили и авторы статьи из Корнелльского университета, которые и описали этот метод. 

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

GPT-Q

Сейчас это далеко не SOTA, но один из самых популярных методов квантизации в индустрии. GPT-Q предполагает хорошую квантизацию только весов в 4 бита, оставляя активации в 16 битах. Должно получиться неплохое ускорение в три раза, как заявляют авторы метода. Но по факту получается только в 1,5–2, причём при условии написания собственной более эффективной реализации CUDA-кернела. Так происходит, потому что production-level фреймворки инференса оптимизированы гораздо лучше, чем те, на которых проводились замеры авторов статьи.

Важно отметить, что ускорение проявляется только в сценарии декодинга, когда мы авторегрессионно генерируем токены. Причём наибольший прирост скорости наблюдается при маленьких батчах — до 8 или 16. 

Наилучшее ускорение наблюдается на больших моделях, поскольку квантизация весов оптимизирует трансфер по памяти.

Мы не будем подробно разбирать устройство метода, а лучше посмотрим бенчмарки. Тут довольно несовременные модели, без LLaMa — только OPT. Квантизация действительно происходит почти без потерь в 4 бита — в 3 бита потери всё-таки есть. В своих замерах мы выяснили, что также без потерь квантизуются современные архитектуры Decoder Only моделей (e.g. LLaMa), даже если внести в них какие-то архитектурные модификации вроде Group Query Attention (GQA) (ускорения от GPT-Q и GQA складываются). Но потери в качестве, исходя из наших экспериментов, немного усугубляются на моделях <= 7B. При этом на моделях размера вроде 35B их нет.  

Впрочем, много кто использует такую модель в трёхбитном формате, благодаря чему с инференсом 175-миллиардной модели может работать одна видеокарта. Это даёт новые возможности тем, у кого есть какие-либо ограничения в ресурсах GPU.

Сильно более подробное введение и обзор методов квантизации LLM можно почитать в другом нашем посте.

Другие идеи и методы

Первый метод — Speculative Decoding. Предположим, что у нас есть большая модель (например, 70B LLaMa) и маленькая (7B). Их можно скомбинировать в одном сервисе. На маленькую возложим генерацию токенов, а на большую — верификацию.  

Мы можем нагенерировать какое-то количество токенов (К в формуле ниже), а затем за один forward pass большой моделью понять, подходят нам они или нет, какую часть токенов принять, а какую выкинуть. Затем можно «бесплатно» сгенерировать ещё один токен большой моделью и запустить процесс заново. Снова генерируем К токенов маленькой моделью, верифицируем и так далее. 

При этом все текущие методы Speculative Decoding оптимизируют одну метрику — Acceptance Rate, то есть количество токенов из маленькой модели, которые приняты большой. Чем оно выше, тем выше ускорение и тем больше мы экономим вызовов большой модели.

Следующий класс техник — Continuous Batching. Рассмотрим пример. На схеме жёлтым цветом отмечены токены контекста, а синим — первый сгенерированный токен. Красным отмечен конец сгенерированного предложения. 

Здесь проблема в возникновении пропусков, или бабблов. Генерируемые гипотезы не являются одинаковыми по длине в токенах, поэтому некоторые из них могут завершиться раньше, чем закончится работа над последними. Из-за этого видеокарта будет простаивать, часть её вычислительных ресурсов окажется без нагрузки.

Чтобы решить эту проблему, нужно заменить «бабблы» на обработку примеров из новых запросов на генерацию. Как это сделать — отдельный и сложный вопрос. Подробнее об этом рассказано в этой статье. 

Последний трюк — Parameter-Efficient Fine-Tuning, или PEFT. Мы уже рассказывали о нём на Хабре. У нас есть одна большая модель, например YandexGPT. К ней можно сделать дообучаемые добавки Prompt Tuning или LoRA. Они весят очень мало, порядка 200 килобайт, и, по сути, это некие картриджи, которые специализируют большую общую модель под конкретную задачу. Такой картридж можно отсылать в API вместе с запросом, что позволяет удобно на одном и том же железе тестировать множество моделей и продуктов на небольшой аудитории или срезе данных. Так можно переиспользовать вычислительные ресурсы на инференсе между несколькими потребителями.

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

• Дистилляция подойдёт в том случае, если у вас есть две модели: большая и маленькая. В идеале использовать PPO для обучения агентов, однако описанные в статье трюки помогут обойтись и без алгоритма.

• Квантизация помогает сократить объёмы потребляемых вычислительных ресурсов и вместе с тем поднять скорость инференса. Для квантизации только весов можно использовать популярный GPT-Q, а вот для квантизации и весов, и активаций пригодится уже SmoothQuant.

• Есть и другие способы ускорить инференс, вроде Speculative Decoding и Continuous Batching, которые можно и нужно применять.

• И наконец, можно комбинировать перечисленные методы для получения наилучшего ускорения.

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

а