Ну и наконец рассмотрим, как нам применить обученный трансформер. И применение уже обученной модели сильно отличается от тренировки, потому что теперь приходится учитывать еще один фактор — задержку. Так что давайте разбираться, как устроен инференс LLM и как его правильно масштабировать.
Основы инференса
Итак, мы натренировали трансофрмер и теперь хотим использовать его по назначению — принимать на вход промпт и генерировать в ответ всякие разные новые последовательности токенов/букв.
Генерация по сути дело простое — модель просто выдает новый токен, руководствуясь формулой log p(next token _i | prev tokens), то есть для всех возможных токенов вычисляются логарифмы вероятностей (на самом деле не совсем вероятностей, но в целом это не так важно), а затем мы сэмплируем из этого распределения и получаем новый токен. Далее мы добавляем новый токен к имеющемуся промпту и повторяем весь процесс заново. То есть трансформер (здесь и далее под термином «трансформер» мы будем иметь в виду декодерную часть классического трансформера) это авторегрессионная модель.
На практике все не так просто, потому что если мы каждый раз будем прогонять промпт через модель, это будет слишком затратно — получим сложность O(n^2) на MLP блоке и O(n^3) на блоке внимания при генерации n токенов.
Как этого избежать? По сути, единственная по‑настоящему новая операция при генерации нового токена это подсчет attention score между текущим токеном и всеми предыдущими. attention score для всех предыдущих токенов не меняются, поэтому мы можем просто их закэшировать. Что мы собственно и делаем — мы сохраняем значения векторов key и value для всех предыдущих токенов в KV‑кэш, и для каждого следующего токена просто берем значение его вектора query_i и вычисляем его произведение со всеми предыдущими key_j, а потом еще и с value_j. Таким образом, сложность генерации n токенов падает до квадратической.
Так, с этим разобрались, теперь рассмотрим режимы инференса. Их всего два:
Prefill. Берем входной промпт, системный промпт и что там еще у нас есть, и обрабатываем вот это вот все в один присест параллельно для всех токенов, далее добавляем новые значения key и value в KV‑кэш.
Generation. Следует после Prefill. Берем KV‑кэш и выходные логиты, сэмплируем очередной токен, пихаем его обратно в трансформер, обновляем KV‑кэш, сэмплируем следующий токен и так далее. И так продолжаем до тех пор пока не наткнемся на <EOS> токен или не достигнем предельной длины сообщения.
И вот в наличии этих двух совершенно разных с вычислительной точки зрения режимов заключается серьезная проблема, поджидающая нас во время инференса трансформера. Но о ней мы поговорим позже, а пока представим еще один параметр, на который нам было почти все равно во время обучения модели, но который становится критически важным во время ее применения.
Латентность
В отличие от обучения, в режиме инференса нам приходится заботится не только о вычислениях и памяти, но и о латентости (latency) — пользователю вряд ли понравится, если чат‑бот будет отвечать на его вопрос через 5 минут, да еще и со скоростью 2 токена в минуту. Как вы могли заметить из предыдущего предложения, латентность может быть двух видов:
TTFT — время до первого токена, то есть время, прошедшее между отправкой промпта в модель и появлением первого токена на выходе, характеризует скорость работы режима Prefill
Per‑token latency — сколько времени уходит на генерацию одного токена, характеризует скорость работы режима Decode
Причем в зависимости от типа применения модели важность этих параметров может меняться. Например в случае оффлайн генерации, когда модель выдает ответ сейчас, а прочитают его когда‑нибудь потом, насчет латентности можно не париться. Далее идет упомянутый выше чат‑бот, для работы которого важно, чтобы TTFT был как можно меньше, а per‑token latency обеспечивал скорость генерации выше скорости чтения пользователя. Ну и наконец агентский сценарий, когда агент должен и начать работу как можно быстрее, и скорость его работы должна быть чем выше, тем лучше.
Однако помимо требованиям к латентности остаются требования к памяти, нагрузке на железо, качеству работы, поэтому инференс это всегда компромисс между этими четырьмя параметрами.
Теперь рассмотрим оба режима применения трансформера поподробнее.
Prefill
Режим Prefill, по сути, не сильно отличается от режима обучения модели (разве что обратного прохода нет) в плане вычислительной нагрузки на ускоритель: у нас есть N токенов из начального промпта, для каждого из этих токенов нужно посчитать внимание, записать соответствующие значения K и V в KV‑кэш, далее прогнать каждый токен через MLP слой. И все это можно делать параллельно — в этом главное сходство с режимом тренировки — нагрузить ускоритель нам достаточно легко.
Как мы уже писали в главе 1, для того, чтобы добиться compute‑bound состояния например на Nvidia H100 для простого матричного умножения (то есть наше оборудование бОльшую часть времени занято нужными нам вычислениями, а не перегонкой данных), размер батча должен быть больше 295 (при условии, что веса и операции хранятся/производятся в одной разрядности) — вполне себе достижимое условие.
Но кое‑что по сравнению с тренировкой все же поменялось. Если во время тренировки модели механизм внимания не играл особого значения (по крайней мере для относительно небольшой длины контекста, на которой модель в основном и тренируется), то в режиме инференса он выходит на первый план. Давайте для начала посчитаем арифметическую интенсивность механизма внимания для разрядности bf16. Итак, что нам нужно сделать для подсчета внимания:
прочитать тензор Q размера bf16[B, T, D] из HBM
прочитать KV‑кэш из HBM, это по сути два тензора K и V размера bf16[B, T, D]
перемножить тензоры Q и K, это потребует 2BTSD FLOPs, ничего обратно в HBM записывать не надо, так как используем Flash Attention
умножить полученный результат на тензор V, на это потребуется 2BTSD операций
записать результат размера bf16[B, T, D] обратно в HBM
B — размер батча, T — длина входной последовательности токенов, S — длина KV‑кэша, D — размерность входных эмбеддингов
Теперь считаем арифметическую интенсивность, которая, напомню, равна отношению общего количества операций к общему количеству переданных данных:
Так как для режима Prefill S = T, то это выражение можно упростить до T^2 / 2T = T / 2. То есть арифметическая интенсивность внимания для режима Prefill равна O(T). А значит в этом режиме ускоритель нагрузить довольно легко — достаточно подать достаточно длинную последовательность токенов.
Decode
Теперь переходим к режиму Decode. С ним все сложнее, так как при генерации каждый следующий токен зависит от предыдущего, а значит каждый запрос может исполняться только в один поток. Эту проблему пытаются решить при помощи speculative decoding или диффузии, но в общем случае будем считать, что генерим по одному токену за раз. А это значит, что для достижения compute bound для MLP слоя в этом режиме нам нужно будет одновременно обслуживать 295 запросов в случае Nvidia H100 — это довольно проблематично.
А что с вниманием? В случае когда T = 1 и S >> T, предыдущее выражение для арифметической интенсивности превращается в S / S = 1. То есть compute bound при вычислении внимания в режиме Decode достичь практически невозможно.
Считаем латентность и пропускную способность LLM
Начнем с латентности. Для небольшого размера батча в режиме генерации мы можем пренебречь вычислениями и просто считать, что ограничены пропускной способностью ускорителя. В таком случае время генерации одного токена будет:
Следовательно пропускная способность (она же скорость генерации токен/с) будет равна 1 / t_token.
Но по мере роста батча вычисления начинают доминировать и тогда время генерации одного токена распадается на сумму времени генерации для механизма внимания и MLP:
Ну и разумеется, между латентностью и пропускной способностью придется искать компромисс. Что наглядно отражено на графике ниже: маленькие батчи работают быстро, но неэффективно, большие — долго, но медленно.
Вывод: если хотим добиться максимальной пропускной способности, надо использовать максимально возможный размер батча данных, если же хотим уменьшить время отклика — размер батча надо уменьшать.
А что с памятью?
А давайте возьмем какую‑нибудь известную нейронку и посчитаем, сколько она кушает памяти. Например LLaMA 2–13 B. Вот ее параметры:
Параметр | Значение |
L (кол‑во слоев) | 40 |
D (размерность входных эмбеддингов) | 5120 |
F (размерность MLP слоя) | 13 824 |
N (кол‑во голов Query) | 40 |
K (кол‑во голов Key и Value) | 40 |
H (размерность голов Q, K, V) | 128 |
V (размер словаря токенов) | 32 000 |
T (длина входной последовательности) | 8192 |
Теперь посчитаем по компонентам:
MLP_params = 3 x D x F x L = 8.5e9 — в каждом MLP слое 3 матрицы
Vocab_params = 2 x V x D = 0.3e9 — один входной слой модели + выходной слой
Attention_params = (2 x D x N x H + 2 x D x K x H) x L = 4.2e9 — тензор Q, тензор O + два тензора K и V
В сумме получается 8.5e9 + 0.3e9 + 4.2e9 = 13e9 или 13B. То есть наш расчет совпадает с ожидаемым.
А теперь посчитаем размер KV‑кэша для одной входной последовательности. Допустим у нас по 2 байта на вес и по 2 байта на активацию, тогда веса нашей модели будут занимать 26 ГБ, а размер KV‑кэша:
KV_size = 2 x 2 x H x K x T x L = 2 x 2 x 128 x 40 x 40 x 8192 = 6.7e9 байт = 6.7 ГБ
То есть KV‑кэш всего для одной последовательности уже занимает 1/4 от веса всей модели. А теперь вспомним, что мы хотели, чтобы наши вычисления в режиме Decode были compute bound хотя бы для MLP слоя. А для этого минимальный размер батча при работе на Nvidia H100 должен быть 295. А значит размер KV‑кэша взлетает до 6.7 x 295 = 1976.5 ГБ, а общий размер потребляемой памяти составит 2 ТБ! И это для дохлой миниатюрной модели. Представьте, сколько будут жрать монстры с триллионами параметров...
Разумеется, в таком виде все это оставлять нельзя, никакой памяти не напасешься, поэтому разработчики придумали всякие разные трюки для того, чтобы ускорить процесс генерации токенов, а также уменьшить задержки и, разумеется, потребление памяти.
Grouped‑Query Attention (GQA)
Идея простая: а давайте просто назначать каждой голове K и V несколько голов Q, это значительно сократит размер KV‑кэша, а на производительность модели при этом повлияет не сильно. Если совсем все плохо с памятью, то можно оставить вообще одну голову K и одну голову V на слой и получить Mulit‑Query Attention.
Рассмотрим наш случай с Llama-2. Если внедрить в нее GQA 1:8 (1 голова K и V на 8 голов Q), то KV‑кэш сократится в 8 раз, примерно до 250 ГБ — многовато, но все же заметно меньше чем 2 ТБ.
Общие KV‑кэши для нескольких слоев
Еще одна простая идея — пусть несколько слоев используют один KV‑кэш. Также сокращает общий размер KV‑кэша, но может получиться так, что эти общие кэши придется считывать по несколько раз из HBM, что не очень хорошо скажется на скорости модели, особенно если ее слои шардированы между ускорителями.
Квантизация
Тут тоже все просто: храним значения KV‑кэша не в bf16, а в fp8, int8, int4 и так далее. При этом, если распределения весов модели находится в достаточно узком диапазоне (а так обычно и бывает), производительность сильно не просядет или не просядет совсем. Дополнительным плюсом квантизации является то, что мы можем тренировать модель в полной разрядности, а квантизировать уже после (Post Training Quantization — PTQ).
Paged Attention
Нам не всегда требуется доступ к полной длине контекста, поэтому имеет смысл хранить KV‑кэш не целиком, а постранично, как это делает ОС. И следовательно добавлять очередную страницу с кэшем только тогда, когда она нужна. Это сильно усложняет саму работу с кэшем, но взамен мы получаем возможность управлять размером доступной нам памяти и не раздувать ее сверх меры.
Масштабирование инференса
Теперь перейдем к тому, зачем мы здесь сегодня собрались — собственно как это все дело масштабировать?
Prefill
В этом режиме все почти также как и в режиме обучения, описанном в предыдущей главе — можно применять любую из перечисленных там техник, главное следить за размером KV‑кэша. И так же как в режиме обучения, увеличение количества ускорителей дает нам бОльшую производительность (снижая таким образом TTFT), но требует дополнительной синхронизации между ускорителями (снижая пропускную способность в пересчете на ускоритель).
Типичный пример шардирования в режиме Prefill:
Сначала шардируем модель, используя Tensor Parallelism: TP режет скрытую размерность MLP слоя, сама ось последовательности при этом остаётся нетронутой, поэтому её можно отдать под отдельную ось шардирования.
Наращиваем TP, пока не упрёмся в пропускную способность соединений между ускорителями — порядка F / (C/W) на одну ось массива ускорителей (F — размерность MLP слоя, С — вычислительная мощность ускорителя, W — пропускная способность шины ускорителя), что для H100 даёт примерно 4–8 шардов.
Дальше подключаем Sequence Parallelism (то есть шардируем вдоль оси последовательности): придется конечно периодически обмениваться кусками KV‑кэша между ускорителямя, но эти коммуникации перекрываются вычислениями и на длинных контекстах стоят немного.
Decode
С генерацией токенов все сложнее. В этом режиме сложно получить большой батч, потому что для этого надо объединить множество запросов, а они обычно разнесены во времени, и пользователи не хотят ждать, пока их запросы объединят и обработают.
Ну и соответственно не все стратегии шардирования тут уместны:
FSDP применить невозможно: нам и так приходится постоянно подгружать веса и KV‑кэш из HBM, еще не хватало гонять их между ускорителями. Нам нужно пересылать активации, не веса.
Data Parallelism тоже не катит: во‑первых, нужно еще постараться впихнуть модель в один ускоритель, а во‑вторых, даже если получится, у нас на каждом ускорителе будет просто копия модели — быстрее, чем справится один ускоритель, нам разогнаться не выйдет.
У нас нет последовательности как таковой, мы генерируем по одному токену за раз — пока, Sequence Sharding.
Pipeline Parallelism — тут все упирается в то, что ускорители расположены последовательно: то есть при обработке токена не получится подгрузить все параметры на всех ускорителях одновременно, придется ждать, когда первый ускоритель передаст результат выполнения второму, затем второй третьему и так далее. В режиме обучения это не проблема, но тут задержка играет ключевую роль, поэтому не вариант.
Так что вариант с Tensor Parallelism, который как раз пересылает активации, будет работать и здесь.
Как работает эффективный движок инференса
Теперь давайте разберемся, как нам делать весь этот инференс эффективно.
Самый простой метод это просто сначала обработать батч в режиме Prefill, затем этот же батч, но в режиме Decode.
Сделать так довольно просто, но есть куча проблем:
В любом случае будет очень большая задержка. Будет большой батч — сможем обрабатывать множество запросов пользователей одновременно, но TTFT будет тоже большим. Маленький батч — уменьшим TTFT, но зато получим очень низкую производительность в режиме генерации. Ну и к тому же перед тем как начать генерацию токенов, нужно будет завершить Prefill для всех запросов — TTFT вырастет еще больше.
Длинные генерации блокируют короткие. Весь процесс не закончится до тех пор, пока не будут обработаны все последовательности, если какие‑то из них закончились раньше — придется ждать пока не закончится самая длинная.
В режиме Prefill все последовательности придется дополнять до длины максимальной последовательности в батче. Это пустая трата вычислений.
Одинаковое шардирование для обоих режимов. А, учитывая особенности этих режимов, перечисленные в предыдущей главе, это мягко говоря не оптимально.
В общем для одного пользователя такой режим подойдет, для множества — едва ли.
Немного другой подход: давайте выполнять Prefill для каждой последовательности отдельно, а затем объединять эти последовательности в батч во время генерации.
Это решает многие проблемы предыдущего подхода, но одна остается — придется каждый раз ждать, пока не закончится очередной Prefill, то есть генерация токенов будет постоянно прерываться, что тоже не очень хорошо с точки зрения пользовательского опыта.
Поэтому приходим к третьему, самому логичному подходу — разделяем Prefill и Decode. Сначала делаем Prefill (один запрос — один сервер), потом посылаем полученные KV‑кэши на серверы для генерации, объединяем их и обрабатываем множество запросов одновременно в режиме Decode.
В этом случае мы и решаем проблему блокировки режимом Prefill режима Decode для разных пользователей, нам легче масштабировать систему в случае изменения соотношения количества вычислений, необходимых для Prefill и Decode, мы можем шардировать ускорители для Prefill и Decode по‑разному, а система «один Prefill запрос — один сервер» избавляет нас от лишних вычислений в случае, когда приходится обрабатывать несколько промптов разной длины.
Появляется другая проблема — нужно гонять KV‑кэши между серверами, но она вполне решаема, к тому же у нас есть эффективные техники для уменьшения размера KV‑кэша.
Continuous batching
Развитие предыдущей идеи:
prefill‑функция добавляет результат Prefill фазы в KV‑буфер с заранее заданной максимальной длиной контекста / количеством страниц, сам результат при этом может иметь разную длину
generate‑функция проходит в кольцевом режиме одновременно токен‑за‑токеном по всем последовательностям в KV‑буфере и генерирует ответ для всех непустых последовательностей
Prefix caching
Начало промпта LLM обычно одинаково для всех запросов — там содержатся всякие инструкции по работе, правила, список инструментов и так далее. А значит эта часть KV‑кэша для всех запросов будет одинакова. Следовательно напрашивается идея — закэшировать KV‑кэш хотя бы частично, это поможет сэкономить на вычислениях. Также можно кэшировать часто попадающиеся вопросы, задачи и так далее
Но поскольку KV‑кэш занимает много памяти, все подряд кэшировать не получится, поэтому для этого используют свободную HBM память или свободную RAM память — хоть она и медленнее HBM, но достаточно быстра для простого чтения из кэша.
На практике:
Раз KV‑кэш лежит локально на той реплике которая обработала исходный запрос, нужна так называемый affinity‑маршрутизация, чтобы запросы от одного пользователя всегда попадали на одну и ту же реплику. А это, в свою очередь, усложняет балансировку нагрузки.
Уменьшение размера KV‑кэша снова актуально: в тот же объём памяти влезет больше кэшей, да и читаются они быстрее.
Для поиска по кэшам хорошо подходит обычное или префиксное деревья (trie). Инвалидацию кэша можно делать с помощью LRU.
Ну и разумеется, все эти навороты уже есть во всех современных фреймворках по работе с LLM (vLLM, SGLang, JetStream и так далее), поэтому самому ничего из этого реализовывать не надо. Красота!
Ну а на этом наш цикл статей о масштабировании LLM подходит к концу. Спасибо всем, кто читал, ставил классы и писал комментарии, пока!
