Алгоритмы спекулятивного инференса LLM

Введение

За последние годы качество LLM моделей сильно выросло, методы квантизации стали лучше, а видеокарты мощнее. Тем не менее качество генерации все еще напрямую зависит от размера весов и, как следствие, вычислительной сложности.
Кроме того, генерация текста авторегрессионна – токен за токеном по одному, потому ее сложность зависит от размера контекста и количества генерируемых токенов.

Но генерация текста не всегда имеет однородную сложность, так же как мы во многом мыслим идеями, а слова произносим “на автомате”. В статье обсудим алгоритмы, позволяющие использовать эту неоднородность для ускорения.

Prefill vs decode

При генерации токенов llm имеет две стадии:

1. Prefill – обработка контекста и заполнение KVCache
   Compute-bound задача (умножение больших матриц)

2. Decode – генерация новых токенов авторегрессионно: сгенерировали токен, добавили к текущему накопленному промпту, повторили, пока не сгенерируем нужное количество
   Memory-bound задача (больше таскаем KVCache и собираем матрицы, чем умножаем)

Тут есть 2 проблемы:

1. При увеличении размера контекста работа с kv-cache становится неэффективной.
   От чего актуальны алгоритмы сжатия и офлоадинга кешей, а также более эффективная работа с памятью ^[1], снижающая фрагментацию. В эту сторону стоит почитать про фреймворк vllm и устройство PagedAttention.
   Я писал об этом в своем телегам канале ^[2] (и о многом другом касаемо инференса нейросетей).

2. При увеличении размера генерации происходит очень много forward вычислений базовой сети, на каждый новый токен, что не эффективно. И это приводит к идее спекуляции.

Идея спекуляции

Имеем по одному forward на каждый новый токен в decode. При этом некоторые слова понятны из контекста или могут не нести смысла в контексте конкретной генерации.

Можем ли мы пропустить или существенно упростить вычисления некоторых токенов?
Ведь если у нас есть “угаданная” последовательность токенов, которая принимается нашей моделью целиком (через сравнение с вероятностями на выходе после общего forward), то можно сократить количество forward итераций на decode стадии.

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

Спекулятивный декодинг

Сгенерируем наперед авторегрессионно токены с помощью драфт модели. После этого за один forward проведем новые сгенерированные токены. Таким образом имеем для каждого токена вероятности q(x) для драфт модели и p(x) для этих же токенов от основной модели.

Теперь, вместо семплирования x ∼ p(x):

1. Семплируем x ∼ q(x)

2. Если q(x) ≤ p(x) принимаем спекулятивные токен

3. Если q(x) > p(x), то реджектим его с вероятностью 1 − p(x)/q(x)

4. В случае реджекта семплируем x из распределения norm(max(0, p(x) − q(x)))

В статьях показано, что такой способ позволяет получить x ∼ p(x).

Fast Inference from Transformers via Speculative Decoding ^[3]
code: https://github.com/feifeibear/LLMSpeculativeSampling ^[4]

Accelerating Large Language Model Decoding with Speculative Sampling ^[5]
code: https://github.com/shreyansh26/Speculative-Sampling ^[6]

Lookahead decoding

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

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

Такая идея предложена в следующей статье. Спекуляция делается через решения системы нелинейных уравнений методом Якоби.

Break the Sequential Dependency of LLM Inference Using LOOKAHEAD DECODING ^[8]
code: https://github.com/hao-ai-lab/LookaheadDecoding ^[9]

Есть реализация в TensorRT-LLM:
https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/lookahead/README.md ^[10]

О самом методе Якоби для параллельной генерации можно почитать тут:

Accelerating Transformer Inference for Translation via Parallel Decoding ^[11]
code: https://github.com/teelinsan/parallel-decoding ^[12]

Архитектурные подмены (early exit и skip decode)

Нужно переобучение с добавлением early-exit слоев или изменением процесса обучения ^[13]. Возможно делать дообучение на готовых весах. Реализация такого подхода не будет быстрой и потребует хорошей компетенции от команды.

SkipDecode: Autoregressive Skip Decoding with Batching and Caching for Efficient LLM Inference ^[14]
code: закрыт 😔

EE-LLM: Large-Scale Training and Inference of Early-Exit Large Language Models with 3D Parallelism ^[15]
code: https://github.com/pan-x-c/EE-LLM ^[16]

Multitoken prediction

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

Better & Faster Large Language Models via Multi-token Prediction ^[17]
code: нет 😔

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

Medusa

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

Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads ^[18]
code: https://github.com/FasterDecoding/Medusa ^[19]

Eagle

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

Работает так же, как и классический спекулятивный декодинг, только модель помимо токенов с прошлых итераций принимает еще и feature-вектора этих токенов (значения выхода после модели и перед LMHead).
Draft модель использует LMHead и слой эмбеддинга токенов с основной модели, обучая только “тушку” на смеси классификационного лосса и l1 (между feature векторами основной и драфт моделей). Манипулируя размером тушки можно достигать оптимального соотношения размера драфт модели и принятия токенов основной моделью. А сам спекулятивный семплинг делается через tree-search (генерируем по 2 токена на каждой итерации).

EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty ^[20]
code: https://github.com/SafeAILab/EAGLE ^[21]

Поддержка фреймворками – TensorRT-LLM (https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/eagle ^[22]), vLLM (https://github.com/vllm-project/vllm/pull/6830 ^[23])

Что использовать?

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

Можно выделить три пути:

1. Lookahead – у вас нет ни драфт модели, ни ресурсов для дообучения,

2. Спекулятивный декодинг – если есть draft модель и можно загружать ее веса вместе с основной моделью. Для большинства открытых нейросетей есть веса меньших размеров, так что метод прост для использования. Например, для предобученной модели 70b можно пробовать драфт модели 1b/3b/7b,

3. Вариации Eagle/Medusa с обучением головы – если есть ресурс на дообучение, доступность и компетенция разработчиков. Можно ускорится относительно спекулятивного декодинга, при этом веса драфт модели будут кратно меньше. Это позволит эффективнее утилизировать память видеокарт.