[Перевод] Как работает инференс в больших языковых моделях

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

Что такое большие языковые модели?

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

Фундаментальная составляющая этих моделей — это трансформерный уровень, состоящий из двух основных компонентов:

1. Механизм самовнимания и 

2. Нейронная сеть прямого распространения.

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

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

Размер модели = количество параметров в сети. Модель на 7 миллиардов параметров содержит 7 миллиардов чисел с плавающей точкой, в которых хранятся знания, усвоенные моделью в процессе обучения. Эти параметры упорядочены в виде матриц весов, которые на каждом уровне преобразуют полученные на вход данные.

Такие модели как GPT-4, Claude и Llama — это трансформеры, выполняющие лишь декодирование. Это означает, что из оригинальной архитектуры трансформеров они используют лишь половину, то есть, декодер. Таким образом, они авторегрессионные, генерируют на каждом проходе новый токен, исходя из всех ранее сгенерированных токенов. Это отлично подходит для решения задач, связанных с генерацией текста.

Токенизация

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

В современных больших языковых моделях наиболее распространён метод токенизации, именуемый Byte Pair Encoding (BPE) — склеивание самых частых пар символов. BPE начинает работу со словаря отдельных символов и далее, итерация за итерацией, объединяет наиболее частые пары смежных токенов, создавая таким образом новые токены.

# Пример процесса токенизации по алгоритму BPE 
# Получен на вход текст: "unhappiness"

# Исходно: ['u', 'n', 'h', 'a', 'p', 'p', 'i', 'n', 'e', 's', 's']
# После слияний: ['un', 'happi', 'ness']

Поскольку здесь используется BPE, распространённые слова представляются в виде отдельных токенов (это эффективно), а редкие или неизвестные слова дробятся на знакомые фрагменты-подслова (это гибко).

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

# Упрощённый процесс токенизации
text = "The AI model generates text"
tokens = tokenizer.encode(text)
# [464, 15592, 2746, 18616, 2420]

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

Векторные представления токенов

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

Если модель оперирует словарём из 50 000 токенов, а измерение векторных представлений равно 4 096, то матрица векторных представлений имеет форму [50000, 4096]. Каждая строка соответствует одному токену, а значения в этой строке образуют векторное представление данного токена.

token_id = 464  # ID токена для слова "The"
embedding_vector = embedding_matrix[token_id]
# Результат: вектор, состоящий из 4096 чисел с плавающей точкой

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

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

Архитектура трансформера

Трансформер обрабатывает векторные представления слой за слоем. На каждом слое трансформера выполняются две основные операции: многоголовое самовнимание и прямое распространение по нейронной сети.

Механизм самовнимания вычисляет для каждого токена три матрицы: запрос (Q), ключ (K) и значение (V). Они получаются путём умножения входных векторных представлений на три изученные матрицы весов.

# Вычисление самовнимания
Q = input @ W_query   # Форма: [batch, seq_len, dim]
K = input @ W_key     # Форма: [batch, seq_len, dim]
V = input @ W_value   # Форма: [batch, seq_len, dim]

Матрицы весов W_query, W_key и W_value усваиваются в ходе обучения. Они инициализируются в случайном порядке, а затем корректируются в процессе обратного распространения так, чтобы из векторных представлений извлекались наиболее полезные паттерны.

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

# Баллы внимания
scores = (Q @ K.transpose()) / sqrt(dim)
attention_weights = softmax(scores)
output = attention_weights @ V

Масштабный коэффициент (квадратный корень из измерения) не допускает, чтобы скалярные произведения принимали слишком большие значения, поскольку в противном случае функция softmax насыщалась бы и выдавала в ходе обучения слишком малые градиенты. 

При многоголовом внимании этот процесс выполняется многократно (распараллеливается) с применением разных изученных матриц проекций. В модели могут использоваться 32 головы внимания, и каждая учится фокусироваться на разных аспектах взаимосвязей между токенами. Значения, полученные на выход от всех голов, далее сцепляются и проецируются обратно на измерение модели.

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

# Сеть прямого распространения
hidden = activation(input @ W1 + b1)   # Расширение с учетверением измерений
output = hidden @ W2 + b2              # Проецирование обратно

Этапы инференса — предзаполнение и декодирование

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

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

# Вычисление на этапе предзаполнения
input_tokens = [token_1, token_2, ..., token_n]
# Обрабатывает все токены одновременно
for layer in model.layers:
    Q, K, V = compute_qkv(input_tokens)
    attention_output = attention(Q, K, V)
    layer_output = feedforward(attention_output)

На этапе предзаполнения модель выдаёт первый результирующий токен и собирает кэш ключей и значений (KV), который мы вскоре обсудим. Показатель «время до первого токена» (TTFT) позволяет измерить длительность этой фазы и напрямую влияет на опыт пользователя при работе с системой — сколько пользователю придётся ждать, пока он увидит какой-либо вывод.

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

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

# Вычисления на этапе декодирования
current_token = first_generated_token
while not done:
    # Вычислять только для нового токена
    q_new = compute_query(current_token)

    # Извлечь кэшированные K, V от предыдущих токенов
    k_cached, v_cached = retrieve_cache()

    # Вычислить внимание на основе кэшированных значений
    attention_output = attention(q_new, k_cached, v_cached)

    next_token = generate_token(attention_output)
    current_token = next_token

Задержка между токенами (ITL) — это время между двумя следующими друг за другом генерациями токенов на этапе декодирования. Этот параметр определяет, как быстро поток текста поступает пользователю после того, как начнётся генерация.

Кэш ключей и значений

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

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

Впечатление о том, как это происходит, можно составить по следующему псевдокоду.

class KVCache:
    def __init__(self):
        self.cache_k = None
        self.cache_v = None
    
    def update(self, new_k, new_v):
        if self.cache_k is None:
            self.cache_k = new_k
            self.cache_v = new_v
        else:
            # Конкатенация новых K, V с кэшированными значениями
            self.cache_k = concat([self.cache_k, new_k], dim=1)
            self.cache_v = concat([self.cache_v, new_v], dim=1)
    
    def get(self):
        return self.cache_k, self.cache_v

Модель ведёт отдельные кэши ключей и значений для каждого уровня трансформера и каждой головы внимания. При генерации n-ного токена в кэше сохраняются матрицы ключей и значений для всех n-1 предыдущих токенов.

Благодаря кэшированию ключей и значений можно драматически ускорить работу. Как показывает практика, при кэшировании ключей и значений можно сгенерировать 1000 токенов примерно за 10 секунд, тогда как без кэширования на ту же задачу затрачивается 50 секунд — пятикратная разница.

Но на ведение кэша ключей и значений тратится память. Кэш растёт линейно вслед за ростом последовательности. При работе с моделью на 13 миллиардов параметров, такой как LLaMA-2, на каждый выходной токен требуется выделить примерно 1 МБ пространства в хранилище кэша. Контекст в 1000 токенов требует 4 ГБ только под кэш, что сравнимо с размером самой модели.

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

Когда я впервые самостоятельно экспериментировал с прогоном модели, мне казалось, что именно из-за графического процессора отклики приходят медленно. Лишь потом я заметил, что кэш ключей и значений постоянно выливается из памяти GPU.

Всякий раз, когда пользователь вводил длинный промпт, задержка резко возрастала. Хватило одного исправления — уменьшить точность с FP16 до INT8 – чтобы сократить время отклика более чем наполовину. 

Перемножение матриц

Перемножение матриц — это вычислительная суть трансформерного инференса. На каждом уровне трансформера матрицы перемножаются многократно: на основе входных значений вычисляются Q, K, V, применяется внимание, эксплуатируется нейронная сеть прямого распространения. 

В графических процессорах для эффективного перемножения матриц применяется стратегия мозаичного деления (tiling strategy). Большая матричная операция делится на более мелкие фрагменты, такие, которые помещаются в разделяемую память. Так сокращается количество дорогостоящих глобальных обращений к памяти.

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

Этот процесс можно ещё сильнее ускорить при помощи тензорных ядер, полностью перенеся на аппаратный уровень целые операции перемножения небольших матриц. В этой модели программирования поддерживаются операции 16x16x16, но на железе они автоматически выполняются как множественные операции 4x4x4.

Точность и квантизация при инференсе

Инференс в больших языковых моделей зачастую выполняется при более низкой точности, чем обучение. Тогда как при обучении обычно применяется точность на уровне FP32 или BF16, при инференсе можно использовать FP16, INT8 или даже INT4, потеря качества будет минимальной.

FP16 (16-разрядные числа с плавающей точкой) позволяют наполовину снизить требования к расходу памяти и полосе передачи данных по сравнению с FP32. Тензорные ядра достигают максимальной пропускной способности при работе с FP16, поэтому такая точность по умолчанию используется во многих конфигурациях, развёрнутых для инференса.

# Форматы точности
FP32: 1 sign bit, 8 exponent bits, 23 mantissa bits
FP16: 1 sign bit, 5 exponent bits, 10 mantissa bits  
BF16: 1 sign bit, 8 exponent bits, 7 mantissa bits
INT8: 8 bits for integer representation
INT4: 4 bits for integer representation

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

Модель на 7 миллиардов параметров, работающая с точностью FP16, требует примерно 14 ГБ памяти (7 миллиардов параметров × 2 байта на параметр). При квантизации до INT4 этот объём снижается до 3,5 ГБ, что позволяет выполнять инференс на обычном железе, которое можно купить в магазине. 

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

Сквозной поток инференса

Шаг 1: Токенизация. Вы пишете промпт “Объясни, как работают трансформеры”, и токенизатор преобразует его в ID токенов. Алгоритм BPE разбивает эти токены на единицы-подслова, и у него получается что-то вроде [Объясни, как, трансформ, еры, работают].

prompt = "Explain how transformers work"
token_ids = tokenizer.encode(prompt)
# [22163, 703, 4659, 364, 990]

Шаг 2: Поиск векторных представлений. ID каждого токена индексируется и попадает в матрицу векторных представлений, по нему извлекается соответствующее ему векторное представление. Если у модели 4096 измерений, то каждый токен превращается в вектор из 4096 чисел с плавающей точкой.

embeddings = embedding_matrix[token_ids]
# Форма: [5, 4096]

Шаг 3: Добавляем позиционное кодирование. Модель добавляет в векторные представления информацию о позициях, так, чтобы механизм внимания знал, в каком порядке идут токены. 

positions = [0, 1, 2, 3, 4]
positional_embeddings = positional_encoding[positions]
input_embeddings = embeddings + positional_embeddings

Шаг 4: Этап предзаполнения. Поступающие на вход векторные представления проходят через трансформер уровень за уровнем. Если в модели 32 уровня, то это происходит 32 раза. 

hidden_states = input_embeddings
for layer in model.layers:
    # Многоголовое самовнимание
    Q = hidden_states @ W_query
    K = hidden_states @ W_key
    V = hidden_states @ W_value
    
    attention_scores = (Q @ K.T) / sqrt(dim)
    attention_probs = softmax(attention_scores)
    attention_output = attention_probs @ V
    
    # Сохраняем ключи и значения в кэше для данного уровня
    kv_cache[layer].update(K, V)
    
    # Остаточное соединение и норма уровня
    hidden_states = layer_norm(hidden_states + attention_output)
    
    # Сеть прямого распространения
    ffn_output = feed_forward(hidden_states)
    
    # Остаточное соединение и норма уровня
    hidden_states = layer_norm(hidden_states + ffn_output)

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

Этап 6: Фаза декодирования. Теперь мы декодируем токены по одному. Для каждого нового токена мы вычисляем свежие значения Q, K, V лишь для данного токена, а для всех предыдущих токенов извлекаем из кэша уже имеющиеся значения.

Этап 7: Детокенизация. Наконец, последовательность ID токенов преобразуется обратно в текст, для этого используется словарь токенизатора.

output_text = tokenizer.decode(generated_tokens)

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

Фреймворки для инференса

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

vLLM реализует PagedAttention для эффективного управления кэшем ключей и значений, а также непрерывный батчинг для обеспечения высокой пропускной способности. Достигаемая в ней пропускная способность в 2-4 раза выше, чем в наивных реализациях на том же самом железе.

TensorRT-LLM от NVIDIA предоставляет сильно оптимизированные ядра, предназначенные именно для графических процессоров NVIDIA, благодаря чему пиковая производительность практически достигает максимума, возможного в теории. Здесь задействуются такие технологии, как объединение в пакеты на лету и поддержка квантизации FP8.

# Упрощённый пример использования vLLM 
from vllm import LLM

llm = LLM(
    model="meta-llama/Llama-2-7b-hf",
    tensor_parallel_size=2,  # Use 2 GPUs
    dtype="float16"
)

outputs = llm.generate(
    prompts=["Explain transformers", "What is AI?"],
    sampling_params=SamplingParams(max_tokens=100)
)

Инференс для генерации текста  (TGI) от Hugging Face обеспечивает широкую поддержку моделей и предоставляет такие возможности как непрерывный батчинг и потоковая передача токенов. Он также предоставляет готовый к использованию в продакшене HTTP API для развёртывания моделей. 

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

Метрики производительности и мониторинг

Чтобы понимать и мониторить инференс, нужно отслеживать несколько ключевых метрик.

Время до первого токена (TTFT) измеряет задержку на этапе предзаполнения. Этот параметр прямо влияет на то, насколько удобно работать с моделью — как долго пользователю придётся ждать до получения первого вывода. Оптимизация TTFT лежит через эффективную обработку промптов, что зачастую предполагает предзаполнение пакетов или приёмы спекулятивного декодирования.

Задержка между токенами (ITL) измеряет время между поступлением данного и следующего токена на этапе декодирования. При низких значениях ITL работа с моделью протекает бесшовно и плавно. Эта метрика сильно зависит от полосы передачи данных в памяти и от эффективности кэша ключей и значений.

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

# Мониторинг производительности
start_time = time.now()
first_token = model.generate_first_token(prompt)
ttft = time.now() - start_time

token_times = []
for i in range(num_tokens):
    token_start = time.now()
    token = model.generate_next_token()
    token_times.append(time.now() - token_start)

itl = mean(token_times)
throughput = num_tokens / sum(token_times)

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

Давление памяти, в особенности, размер кэша ключей и значений, влияет на максимальную длину контекстного окна и на размер пакета. Отслеживая кэш-память, мы предотвращаем ошибки, связанные с исчерпанием памяти, а также сами руководим решениями о квантизации.

Примечание

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

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

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

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