3 фишки для ускорения LLM

Введение

LLM, или Large Language Models — этот термин плотно вошёл в современную повестку инженерных специалистов, принёс много value для бизнеса и фана для пользователей ChatGPT. Однако, как видно из самого термина, это большие модели, а значит, они требуют больших вычислительных ресурсов. В новой статье мы поговорим про то, как же их сократить и сделать весь процесс чуть быстрее 🙂

Фишка 1: инференс-фреймворки

Если речь заходит про инференс, например, на GPU в компьютерном зрении — опытный специалист обычно советует выбирать инференс-фреймворк, например, TesnorRT или OnnxRuntime. С LLM ситуация несколько отличается: инференс-фреймворков для них становится всё больше и больше. Здесь есть ряд основных игроков (в скобках указано число звёзд у соответствующего репозитория на Гитхабе):›

1. vLLM (≥ 47k ⭐);
2. TensorRT-LLM (≥ 10к ⭐);
3. llama.cpp (≥ 80k ⭐);
4. SGLang (≥ 13k ⭐);
5. LMDeploy (≥ 5к ⭐);
6. Старый добрый TGI от Hugging Face (≥ 10k ⭐);
7. MLC LLM (≥ 20k ⭐);
8. ExLlamaV2 (≥ 4k ⭐).

Какой же фреймворк лучше выбрать?

Простой ответ: vLLM.

Качественно сделанный фреймворк. Популярный, активно развивается, постоянно добавляются новые фичи, есть удобная документация и интерфейс. Предназначен для инференса на GPU. Однако недавно появилась возможность запуска и на CPU (но для этого требуется сборка vLLM из исходников, а не просто «pip install»). Наблюдение: во всех случаях, кроме mixed-precision квантизации (например, 2.71 битной) и запуска на самом новом железе, vLLM либо уже лучший, либо в итоге догонит.

Ответ посложнее: зависит от ситуации (от железа, требований, задачи).

Например, TensorRT-LLM (TRT-LLM) подойдёт на случай, если вам нужна лучшая скорость на самых современных GPU (H200, B200 и т.д.), или вы хотите использовать FP4-квантизацию и при этом готовы терпеть документацию похуже и медленное добавление «новых фич».

Или llama.cpp. Если вы собираетесь запускать модели на своём ноутбуке на CPU. Или хотите использовать на максимум возможности GGUF-квантизации (которая mixed-precision), в том числе и эффективно инферить такие модели. Можно запускаться и на GPU. Кроме того, там поддерживается сценарий, когда часть модели находится на CPU, а часть — на GPU (так можно «умещать» большие модели на маленьких GPU).

Показатели эффективности инференса LLM

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

1. Обеспечить высокую скорость обработки запросов в системах, где необходимо одновременное обслуживание множества пользователей. Оно может возникнуть при создании массового продукта, например, онлайн-переводчика. Чтобы оценить, насколько эффективно система справляется с этой задачей, используется метрика Throughput (Tokens/s), или пропускная способность. Она показывает, сколько токенов система может сгенерировать за единицу времени.
2. Минимизировать задержку перед началом ответа для комфортного взаимодействия пользователя с системой. Оно возникает практически в любых приложениях, связанных с генерацией текста, так как улучшает пользовательский опыт. Для измерения качества решения задачи используется метрика Time-to-First-Token (TTFT, ms), или время до генерации первого токена. Она показывает, сколько времени проходит с момента отправки запроса до момента, когда система выдаёт первый токен.
3. Обеспечить плавную и быструю генерацию текста после начала ответа, чтобы пользователь не ощущал задержек при выводе информации. Как и предыдущее требование, плавная генерация улучшает пользовательский опыт в интерактивных приложениях. Для оценки этой задачи применяется метрика Time-Per-Output-Token (TPOT, ms), или среднее время генерации каждого последующего токена после первого. Она измеряет, сколько времени нужно системе для генерации каждого нового токена. На рисунке ниже можно увидеть, где во время генерации возникают последние две метрики:

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

Исследование 1: vLLM vs TensorRT-LLM

Сравнение, проведённое ребятами из SqueezeBits:

Мы видим, что в проведённой серии экспериментов TensorRT-LLM везде показал себя лучше.

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

При указанном ограничении в 20 мс пропускная способность у TensorRT-LLM уже получается ниже, чем у vLLM, — вместе с предыдущим это определяет vLLM как наиболее предпочтительное решение для данного кейса.

Исследование 2: vLLM, LMDeploy, MLC LLM, TensorRT-LLM и TGI

Сравнение, проведённое в BentoML.

Оценивали TTFT и Throughput для модели LLama 3 8B (а также Llama 3 70B — с весами, квантованными в 4 бита), запущенной разными фреймворками на карте A100 80GB GPU при разном числе одновременных запросов к модели.

Вывод: в этом эксперименте по TTFT-метрике побеждает vLLM, однако для оптимального Throughput можно присмотреться и к его конкурентам.

Промежуточные выводы

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

Если же вы не знаете, с чего начать, и хотите использовать надёжное решение с новыми фичами для сервера, а заморачиваться — не ваш выбор, то vLLM — то, что надо. Для простых «работяг», которые желают запустить на своём ноуте DeepSeek, чтобы не сливать данные в Интернет, — lama.cpp. А для любителей хардкора — TRT-LLM 😎

Фишка 2: спекулятивное декодирование

Известно, что инференс LLM можно разделить на 2 части: prefill-фаза (то есть обработка входного запроса для генерации самого первого токена) и авторегрессионная генерация (процесс, при котором выход модели подается на её вход итеративно).

Исследователи выяснили, что prefill-фаза занимает значительно меньше времени по сравнению с генерацией.

Можно ли ускорить авторегрессионную генерацию? Оказывается, что да. Например, следующим образом:

1. Взять в дополнение к нашей целевой модели более маленькую (пусть финальное распределение вероятностей на выходе у них p и q);
2. Авторегрессионно сгенерировать текст маленькой моделью, а вот проверить и исправить сгенерированное при помощи нашей целевой модели;
3. Целевая модель при этом работает в prefill-режиме, то есть может обрабатывать сгенерированные токены параллельно;
4. Принять наш токен, если \(\frac{p(x)}{q(x)} \geq 1\), и отвергнуть с вероятностью \(1 — \frac{p(x)}{q(x)}\);
5. «Плохой» токен засемплировать заново с поправленного распределения: \(p'(x) = norm(max(0, p(x) — q(x)))\).

Таким образом, мы сможем ускорить инференс примерно в 2-3 раза по сравнению с базовой моделью.

Мы рассказали про довольно простой способ спекулятивного декодирования, но есть и другие — EAGLE и RecurrentDrafter. Они не требуют дополнительной модели. Подробнее про них и другие более новые методы мы поговорим на нашем курсе «Ускорение нейронных сетей» 😎

Фишка 3: квантование

Если вы читали нашу прошлую статью про квантование LLM, то уже наверняка имеете представление о том, что это важная составляющая инференса модели. Мы не отказываемся от своих слов, поэтому включаем данный пункт и в эту статью 🙂

Общая теория

В зависимости от распределения чисел в квантуемом тензоре можно применять симметричную (используется только скейл) или асимметричную (скейл и зеро поинт) квантизацию.

Квантовать можно как только веса модели (weight-only quantization), так и веса вместе с промежуточными представлениями, или «активациями» (weight – activation quantization). Квантование весов позволяет гарантированно уменьшить размер модели. При дополнительном же квантовании активаций можно рассчитывать на большую скорость инференса (за счёт использования целочисленных вычислений). Однако должна быть программная и аппаратная поддержка таких операций («кернелы»). Но и при weight-only квантизации ускорение тоже может быть. Благодаря, например, Marlin кернелам (которые уже автоматически используются в некоторых LLM инференс-фреймворках, например, в том же vLLM).

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

Квантовать можно весь тензор целиком с одними параметрами квантизации. А можно разбить тензор на несколько частей и квантовать их по отдельности. Таким образом, можно выбирать «гранулярность» квантизации. Так, активации можно квантовать per-token, а веса — per-channel.

Обучать квантованную модель можно параллельно с обычным обучением. Это QAT-квантизация (Quantization aware training). А можно сначала обучить full-precision модель, а потом в конце её отдельно квантануть (либо всё-таки с использованием каких-то калибровочных данных, либо вообще без данных). Это PTQ-квантизация (Post-training quantization).

QAT позволяет получать просадку меньше, чем у PTQ. Но и проводить QAT дольше и сложнее. А для больших языковых моделей QAT-квантизация «отпадает» как совершенно непрактичная.

«Рабочие лошадки» для квантования LLM

Приведём самые популярные методы квантизации LLM и библиотеки, где их можно использовать:

1. LLM.int8() (W8A8) — bitsandbytes от Hugging Face;
2. GPTQ (W4A16) — AutoGPTQ (deprecated), GPTQModel, LLM Compressor;
3. SmoothQuant (W8A8) — LLM Compressor;
4. AWQ (W4A16) — AutoAWQ (deprecated), LLM Compressor;
5. FP8 (W8A8) — LLM Compressor;
6. GGUF (mixed-precision) — llama.cpp (vLLM может такое инферить, но пока не очень хорошо).

Также отметим:

• SmoothQuant может применяться «в связке» с другими методами квантизации (например, GPTQ).
• FP8 лучше, чем SmoothQuant, хотя бы потому что там можно использовать per-tensor режим квантизации (один скейл на тензор, а не построчные скейлы), что даёт прирост в скорости до 10% на больших моделях относительно SmoothQuant.
• В рамках GGUF есть куча вариантов квантизации (и возможность выбирать tradeoff между скоростью / размером и качеством).

Новости

Приведём интересные, на наш взгляд, события, за последнее время:

1. Neural Magic (создатели GPTQ) окончательно подружились с vLLM и закрыли свой форк репоса. Следовательно, в vLLM могут завезти что-то интересное намного раньше, чем в TRT-LLM (ещё раньше, чем уже было до этого).
2. Народ достаточно активно переходит на вычисления в fp8. Более того, в vLLM завезли квантование KV-кэша в этом же типе данных, снизили примерно в 2 раза требования по памяти, что увеличило, в свою очередь, пропускную способность.
3. TRT-LLM окончательно «заопенсорсились», что может значительно увеличить скорость появления новых прикольных фичей.

Flash-attention

Итак, мы разобрали три фишки — инференс-фреймворки, спекулятивное декодирование и квантование. Не забудем и про «базу» — вспомним классические способы оптимизации 🙂

В своё время эта техника шумела: мы получаем ускорение трансформеров, просто эффективно реализуя механизм внимания на низком уровне!

Если коротко: авторы увидели, на что тратится время при работе механизма внимания — на обращение к памяти, а не на сами вычисления. Кроме того, есть разные «виды» памяти c различной пропускной способностью:

Рассмотрим рисунок выше: в основании пирамиды находится DRAM (Direct Random Access Memory), или оперативная память. Её много, но она медленная. В середине лежит HBM (High Bandwidth Memory). Её уже заметно меньше, но и скорость у неё повыше. А на вершине «пищевой цепочки» — SRAM (Static Random Access Memory). Она самая быстрая и самая маленькая.

У flash-attention есть 3 компоненты успеха: тайловый подсчёт механизма внимания, реализация в виде одного CUDA-ядра и пересчёт значений при обратном прогоне. Всё это снижает количество обращений к HBM, что ускоряет инференс.

Тайловый («кусочный») подсчёт механизма внимания

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

Тайловое матричное перемножение экономит обращение к памяти, но для наивного подсчёта softmax нужно снова подгрузить всю строку целиком, что сводит все усилия на нет:

\(m(x) := \max\limits_{i} x_i\)

\(f(x) := \left[ e^{x_1 — m(x)}, \ldots, e^{x_n — m(x)} \right]\)

\(l(x) := \sum\limits_{i} f(x)_i\)

\(\operatorname{softmax}(x) := \frac{f(x)}{\ell(x)}\)

Есть ли способ справиться с этим? Да! Воспользоваться блочным подсчётом softmax:

\(m(x) = m\left([x^{(1)} x^{(2)}]\right) = \max(m(x^{(1)}), m(x^{(2)}))\)

\(f(x) = [e^{m(x^{(1)}) — m(x)} f(x^{(1)}),\ e^{m(x^{(2)}) — m(x)} f(x^{(2)})]\)

\(l(x) = l([x^{(1)} \ x^{(2)}]) = e^{m(x^{(1)}) — m(x)} l(x^{(1)}) + e^{m(x^{(2)}) — m(x)} l(x^{(2)})\)

\(\operatorname{softmax}(x) = \frac{f(x)}{l(x)}\)

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

Пусть у нас есть вектор \(x\) из 4 значений: \(x = [a, b, c, d]\). Между ними отношение порядка: \(a > c > b > d\).

Тогда: \(\operatorname{softmax}(x) = [\frac{e^a}{e^a + e^{b}+e^{c}+e^d},…, \frac{e^d}{e^a + e^{b}+e^{c}+e^d}]\).

Посчитаем \(\operatorname{softmax}\) блочным способом. Разобьём \(x\) на две части: \(x = [x_{1}, x_{2}]\).

Для \(x^{(1)}\): a > b, значит, \(m(x^{(1)}) = a =: m_{1}\).

Для \(x^{(2)}\): c > d, значит, \(m(x^{(2)}) = c =: m_{2}\).

\(f(x^{(1)}) = [e^{a-m_{1}}, e^{b-m_{1}}],\ f(x^{(2)}) = [e^{c-m_{2}}, e^{d-m_{2}}]\) .

\(l(x^{(1)}) = e^{a-m_{1}}+e^{b-m_{1}},\ l(x^{(2)}) = e^{c-m_{2}}+e^{d-m_{2}}\) .

Затем посчитаем для всего вектора \(x\):

\(m(x) = \max(m_{1}, m_2) = a = m\)

\(\begin{split} l(x) &= l(x_1)\,e^{m_1-m} + l(x_2)\,e^{m_2-m}\\ &= (e^{a-m_1}+ e^{b-m_1})\,e^{m_1-m} + (e^{c-m_2}+ e^{d-m_2})\,e^{m_2-m}\\ &= e^{a-m}+ e^{b-m} + e^{c-m}+ e^{d-m} \end{split}\)

Итак, мы получили нормализующий коэффициент. Теперь давайте посчитаем, что у нас в числителе:

\(\begin{split} f(x) &= [f(x_1)\,e^{m_1-m},\ f(x_2)\,e^{m_2-m}]\\ &= \bigl[[e^{a-m_{1}}, e^{b-m_{1}}]\,e^{m_1-m},\ [e^{c-m_{2}}, e^{d-m_{2}}]\,e^{m_2-m}\bigr]\\ &= [e^{a-m}, e^{b-m}, e^{c-m}, e^{d-m}]
\end{split}\)

Далее поделим одно на другое и получим классический \(\operatorname{softmax}\):

\(\begin{split} \operatorname{softmax}(x) &= \frac{f(x)}{l(x)}\\ &= \left[\frac{e^{a-m}}{e^{a-m}+ e^{b-m} + e^{c-m}+ e^{d-m}},\ \ldots,\ \frac{e^{d-m}}{e^{a-m}+ e^{b-m} + e^{c-m}+ e^{d-m}}\right]\\ &= \left[\frac{e^{a}}{e^{a}+ e^{b} + e^{c}+ e^{d}},\ \ldots,\ \frac{e^{d}}{e^{a}+ e^{b} + e^{c}+ e^{d}}\right] \end{split}\)

Может возникнуть вопрос: «А чего здесь сложного? Ведь я могу для \(\operatorname{softmax}\)’а просто «копить» знаменатель, а числитель для каждой компоненты вектора считается независимо». Казалось бы, делю одно на другое и всё? Но дело в том, что при таком подсчёте могут возникнуть переполнения в результате суммирования в знаменателе. А при использовании подсчёта \(\operatorname{softmax}\)’а c помощью формул выше вы сможете обеспечить стабильность вычислений.

Реализация в виде цельного CUDA-ядра

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

Пересчёт значений при обратном прогоне

Для инференса это не особо актуально, но информация довольно полезная. Есть такой подход, как gradient checkpointing. Основная идея следующая: экономить память, не храня промежуточные значения слоёв, необходимые для обратного прогона, и пересчитывать их, когда потребуется. Тут почти то же самое, только на уровне механизма внимания.

Использование flash-attention-v1 (и последующих его версий) — скорее де-факто стандарт, чем что-то необычное. И если вы ещё кипятите используете базовую реализацию, лучше обновить PyTorch 🙂

Заключение

Итак, сегодня мы рассмотрели 3 популярных метода ускорения LLM — инференс-фреймворки, спекулятивное декодирование и квантование.

Надеемся, они помогут вам сделать ваши чат-боты немного быстрее 🙂

References

1. Статья “[vLLM vs TensorRT-LLM] #1. An Overall Evaluation”.
2. Статья “Benchmarking LLM Inference Backends: vLLM, LMDeploy, MLC-LLM, TensorRT-LLM, and TGI”.
3. Статья “A Guide to LLM Inference Performance Monitoring”.
4. Статья «FP32, FP16, BF16 и FP8 — разбираемся в основных типах чисел с плавающей запятой”.
5. Статья “How to tile matrix multiplication”.