Skip to content

Instantly share code, notes, and snippets.

@delfer
Created August 26, 2025 20:21
Show Gist options
  • Save delfer/fca4cd571044d98748e719dbe77af55b to your computer and use it in GitHub Desktop.
Save delfer/fca4cd571044d98748e719dbe77af55b to your computer and use it in GitHub Desktop.

Как работает вывод LLM с помощью llama.cpp

Опубликовано:11 ноября 2023 г. | в 16:00 (34 минуты чтения)

В этой статье мы рассмотрим внутреннюю структуру больших языковых моделей (БЯМ), чтобы получить практическое представление о том, как они работают. Для этого мы будем использовать исходный код llama.cpp, реализацию модели LLaMA от Meta на чистом C++. Лично я считаю, что llama.cpp — отличное учебное пособие для более глубокого понимания БЯМ. Его код чистый, лаконичный и понятный, без лишних абстракций. Мы будем использовать эту версию коммита.

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

Этот пост написан для инженеров в областях, отличных от ML и AI, которые заинтересованы в лучшем понимании LLM. Он фокусируется на внутренней части LLM с инженерной точки зрения, а не с точки зрения AI. Следовательно, он не предполагает обширных знаний в области математики или глубокого обучения.

В этой статье мы рассмотрим процесс логического вывода от начала до конца, затронув следующие темы (нажмите, чтобы перейти к соответствующему разделу):

  • Тензоры: базовый обзор того, как выполняются математические операции с использованием тензоров, которые потенциально могут быть перенесены на графический процессор.
  • Токенизация: процесс разделения запроса пользователя на список токенов, которые LLM использует в качестве входных данных.
  • Встраивание: процесс преобразования токенов в векторное представление.
  • Трансформер: центральная часть архитектуры больших языковых моделей, отвечающая за сам процесс логического вывода. Мы сосредоточимся на механизме самовнимания.
  • Сэмплирование: процесс выбора следующего прогнозируемого токена. Мы рассмотрим два метода сэмплирования.
  • Кэш KV: распространённый метод оптимизации, используемый для ускорения логического вывода в больших подсказках. Мы рассмотрим базовую реализацию кэша KV.

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

Содержание

Откройте оглавление

Высокоуровневый поток данных от запроса к результату

Будучи большой языковой моделью, LLaMA работает следующим образом: на вход подаётся текст, «подсказка», и система предсказывает, какими должны быть следующие токены или слова.

Чтобы проиллюстрировать это, мы возьмём в качестве примера первое предложение из статьи Википедии о квантовой механике. Наша подсказка:

Quantum mechanics is a fundamental theory in physics that

Языковая модель пытается продолжить предложение в соответствии с тем, что, по её мнению, является наиболее вероятным продолжением. Используя llama.cpp, мы получаем следующее продолжение:

provides insights into how matter and energy behave at the atomic scale.

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

Высокоуровневый поток данных от запроса к результату

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

Согласно схеме, процесс выглядит следующим образом:

  1. Токенизатор разбивает запрос на список токенов. Некоторые слова могут быть разбиты на несколько токенов в зависимости от словаря модели. Каждый токен представлен уникальным номером.
  2. Каждый числовой токен преобразуется в эмбеддинг. Эмбеддинг — это вектор фиксированного размера, который представляет токен в более удобной для обработки LLM форме. Все эмбеддинги вместе образуют матрицу эмбеддингов.
  3. Матрица встраивания служит входными данными для трансформера. Трансформер — это нейронная сеть, которая является ядром большой языковой модели. Трансформер состоит из нескольких слоёв. Каждый слой принимает входную матрицу и выполняет над ней различные математические операции с использованием параметров модели, наиболее примечательным из которых является механизм самовнимания. Выходные данные слоя используются в качестве входных данных для следующего слоя.
  4. Последняя нейронная сеть преобразует выходные данные трансформера в логиты. Каждому возможному следующему токену соответствует логит, который представляет собой вероятность того, что токен является «правильным» продолжением предложения.
  5. Для выбора следующего токена из списка логитов используется один из нескольких методов выбора.
  6. Выбранный токен возвращается в качестве результата. Чтобы продолжить генерацию токенов, выбранный токен добавляется к списку токенов, полученному на шаге (1), и процесс повторяется. Это можно продолжать до тех пор, пока не будет сгенерировано нужное количество токенов или пока LLM не выдаст специальный токен конца потока (EOS).

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

Знакомство с тензорами с помощью ggml

Тензоры — это основная структура данных, используемая для выполнения математических операций в нейронных сетях. В llama.cpp используется ggml — реализация тензоров на чистом C++, эквивалентная PyTorch или Tensorflow в экосистеме Python. Мы будем использовать ggml, чтобы понять, как работают тензоры.

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

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

Базовая структура тензора

В ggml тензоры представлены структурой ggml_tensor В упрощённом для наших целей виде это выглядит следующим образом:

// ggml.h
struct ggml_tensor {
    enum ggml_type    type;
    enum ggml_backend backend;

    int     n_dims;
    // number of elements
    int64_t ne[GGML_MAX_DIMS];
    // stride in bytes
    size_t  nb[GGML_MAX_DIMS];

    enum ggml_op op;

    struct ggml_tensor * src[GGML_MAX_SRC];

    void * data;

    char name[GGML_MAX_NAME];
};

Первые несколько полей просты в использовании:

  • type содержит примитивный тип элементов тензора. Например, GGML_TYPE_F32 означает, что каждый элемент представляет собой 32-битное число с плавающей запятой.
  • enum содержит информацию о том, поддерживается ли тензор центральным процессором или графическим процессором. Мы вернёмся к этому позже.
  • n_dims Это количество измерений, которое может варьироваться от 1 до 4.
  • ne содержит количество элементов в каждом измерении. ggml — это порядок, в котором обрабатываются строки, то есть ne[0] обозначает размер каждой строки, ne[1] — размер каждого столбца и так далее.

nb Это немного сложнее. Он содержит шаг: количество байтов между последовательными элементами в каждом измерении. В первом измерении это будет размер примитивного элемента. Во втором измерении это будет размер строки, умноженный на размер элемента, и так далее. Например, для тензора 4x3x2:

Тензорный шаг

Пример тензора из 32-битных чисел с плавающей запятой с размерами {4,3,2} и шагами {4,16,48}.

Цель использования шага состоит в том, чтобы можно было выполнять определённые операции с тензорами без копирования данных. Например, операцию транспонирования двумерного массива, которая преобразует строки в столбцы, можно выполнить, просто поменяв местами ne и nb и указав на те же базовые данные:

// ggml.c (the function was slightly simplified).
struct ggml_tensor * ggml_transpose(
        struct ggml_context * ctx,
        struct ggml_tensor  * a) {
    // Initialize `result` to point to the same data as `a`
    struct ggml_tensor * result = ggml_view_tensor(ctx, a);

    result->ne[0] = a->ne[1];
    result->ne[1] = a->ne[0];

    result->nb[0] = a->nb[1];
    result->nb[1] = a->nb[0];

    result->op   = GGML_OP_TRANSPOSE;
    result->src[0] = a;

    return result;
}

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

Тензорные операции и представления

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

  • op Это может быть любая поддерживаемая операция между тензорами. Установка значения GGML_OP_NONE означает, что тензор содержит данные. Другие значения могут обозначать операцию. Например, GGML_OP_MUL_MAT означает, что этот тензор не содержит данных, а лишь представляет результат матричного умножения двух других тензоров.
  • src — это массив указателей на тензоры, между которыми должна быть выполнена операция. Например, если op == GGML_OP_MUL_MAT, то src будет содержать указатели на два тензора, которые нужно перемножить. Если op == GGML_OP_NONE, то src будет пустым.
  • data указывает на фактические данные тензора или NULL если этот тензор является результатом операции. Он также может указывать на данные другого тензора, и тогда он называется представлением. Например, в приведённой выше функции ggml_transpose() результирующий тензор является представлением исходного тензора, только с изменёнными размерами и шагами. data указывает на то же место в памяти.

Функция умножения матриц хорошо иллюстрирует эти понятия:

// ggml.c (simplified and commented)
struct ggml_tensor * ggml_mul_mat(
        struct ggml_context * ctx,
        struct ggml_tensor  * a,
        struct ggml_tensor  * b) {
    // Check that the tensors' dimensions permit matrix multiplication.
    GGML_ASSERT(ggml_can_mul_mat(a, b));

    // Set the new tensor's dimensions
    // according to matrix multiplication rules.
    const int64_t ne[4] = { a->ne[1], b->ne[1], b->ne[2], b->ne[3] };
    // Allocate a new ggml_tensor.
    // No data is actually allocated except the wrapper struct.
    struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, MAX(a->n_dims, b->n_dims), ne);

    // Set the operation and sources.
    result->op   = GGML_OP_MUL_MAT;
    result->src[0] = a;
    result->src[1] = b;

    return result;
}

В приведённой выше функции result не содержит никаких данных. Это просто представление теоретического результата умножения a на b.

Вычисление тензоров

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

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

// llama.cpp
static struct ggml_cgraph * llm_build_llama(/* ... */) {
    // ...

    // K,Q,V are tensors initialized earlier
    struct ggml_tensor * KQ = ggml_mul_mat(ctx0, K, Q);
    // KQ_scale is a single-number tensor initialized earlier.
    struct ggml_tensor * KQ_scaled = ggml_scale_inplace(ctx0, KQ, KQ_scale);
    struct ggml_tensor * KQ_masked = ggml_diag_mask_inf_inplace(ctx0, KQ_scaled, n_past);
    struct ggml_tensor * KQ_soft_max = ggml_soft_max_inplace(ctx0, KQ_masked);
    struct ggml_tensor * KQV = ggml_mul_mat(ctx0, V, KQ_soft_max);

    // ...
}

Код представляет собой серию тензорных операций и создаёт граф вычислений, идентичный тому, что описан в оригинальной статье о трансформере: Граф вычислений механизма внимания

Чтобы фактически вычислить результирующий тензор (здесь он KQV), необходимо выполнить следующие действия:

  1. Данные загружаются в data-указатель каждого конечного тензора. В примере конечными тензорами являются K, Q и V.
  2. Выходной тензор (KQV) преобразуется в вычислительный граф с помощью ggml_build_forward(). Эта функция относительно проста и упорядочивает узлы в порядке глубины.1
  3. Граф вычислений запускается с помощью ggml_graph_compute(), который выполняет ggml_compute_forward() для каждого узла в порядке обхода в глубину. ggml_compute_forward() выполняет основную часть вычислений. Он выполняет математическую операцию и заполняет указатель data тензора результатом.
  4. В конце этого процесса указатель data выходного тензора указывает на конечный результат.

Передача вычислений на графический процессор

Многие тензорные операции, такие как сложение и умножение матриц, могут выполняться на графическом процессоре гораздо эффективнее благодаря высокому уровню параллелизма. Если графический процессор доступен, тензоры можно пометить как tensor->backend = GGML_BACKEND_GPU. В этом случае ggml_compute_forward() попытается перенести вычисления на графический процессор. Графический процессор выполнит тензорную операцию, и результат будет сохранён в памяти графического процессора (а не в указателе data).

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

Разгрузка графического процессора от тензорных вычислений

Процесс начинается с копирования K``Q, V в память графического процессора. Затем центральный процессор выполняет вычисления тензор за тензором, но фактическая математическая операция выполняется графическим процессором. Когда последняя операция в графе завершается, данные результирующего тензора копируются из памяти графического процессора в память центрального процессора.

Примечание: в реальном трансформере K,Q,V не являются фиксированными, а KQV не является конечным результатом. Подробнее об этом позже.

Теперь, когда мы разобрались с тензорами, можно вернуться к работе LLaMA.

Токенизация

Первым шагом в логическом выводе является токенизация. Токенизация — это процесс разделения запроса на список более коротких строк, называемых токенами. Токены должны входить в словарь модели, то есть в список токенов, на которых обучалась языковая модель. Например, словарь LLaMA состоит из 32 000 токенов и является частью модели.

В нашем примере токенизация разбивает запрос на одиннадцать токенов (пробелы заменяются специальным метасимволом ’ ’ (U+2581)):

|Quant|um|▁mechan|ics|▁is|▁a|▁fundamental|▁theory|▁in|▁physics|▁that|

Для токенизации LLaMA использует токенизатор SentencePiece с алгоритмом кодирования байтовых пар (BPE). Этот токенизатор интересен тем, что он основан на подсчётах слов, то есть слова могут быть представлены несколькими токенами. Например, в нашем запросе слово «Quantum» разбивается на «Quant» и «um». Во время обучения, когда формируется словарь, алгоритм BPE гарантирует, что распространённые слова будут включены в словарь как единый токен, а редкие слова будут разбиты на подсчёты слов. В приведённом выше примере слово «Quantum» не входит в словарь, а «Quant» и «um» являются двумя отдельными токенами. Пробелы не обрабатываются отдельно и включаются в токены в качестве метасимволов, если они достаточно распространены.

Токенизация на основе подслов эффективна по нескольким причинам:

  • Это позволяет магистру права выучить значение редких слов, таких как "Quantum", сохраняя при этом относительно небольшой объем словаря, представляя распространенные суффиксы и приставки в виде отдельных лексем.

  • Он изучает особенности конкретного языка, не используя схемы токенизации, специфичные для этого языка. Цитата из статьи о BPE-кодировании:

    Рассмотрим такие словосочетания, как немецкое Abwasser|behandlungs|anlange «станция очистки сточных вод», для которых сегментированное представление переменной длины интуитивно более привлекательно, чем кодирование слова в виде вектора фиксированной длины.

  • Аналогичным образом он полезен при разборе кода. Например, переменная с именем model_size будет преобразована в model|_|size, что позволит языковой модели «понять» назначение переменной (ещё одна причина давать переменным информативные имена!).

В llama.cpp токенизация выполняется с помощью функции llama_tokenize() . Эта функция принимает на вход строку с подсказкой и возвращает список токенов, где каждый токен представлен целым числом:

// llama.h
typedef int llama_token;

// common.h
std::vector<llama_token> llama_tokenize(
        struct llama_context * ctx,
        // the prompt
        const std::string & text,
        bool   add_bos);

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

Q|u|a|n|t|u|m|▁|m|e|c|h|a|n|i|c|s|▁|i|s|▁a|▁|f|u|n|d|a|m|e|n|t|a|l|

Qu|an|t|um|▁m|e|ch|an|ic|s|▁|is|▁a|▁f|u|nd|am|en|t|al|

Qu|ant|um|▁me|chan|ics|▁is|▁a|▁f|und|am|ent|al|

Quant|um|▁mechan|ics|▁is|▁a|▁fund|ament|al|

Quant|um|▁mechan|ics|▁is|▁a|▁fund|amental|

Quant|um|▁mechan|ics|▁is|▁a|▁fundamental|

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

Встраивания

Токены используются в качестве входных данных для LLaMA, чтобы предсказать следующий токен. Ключевой функцией здесь является функция llm_build_llama():

// llama.cpp (simplified)
static struct ggml_cgraph * llm_build_llama(
         llama_context & lctx,
     const llama_token * tokens,
                   int   n_tokens,
                   int   n_past);

Эта функция принимает на вход список токенов, представленных параметрами tokens и n_tokens. Затем она строит полный граф тензорных вычислений LLaMA и возвращает его в виде struct ggml_cgraph. На этом этапе вычисления не производятся. Параметр n_past, который в настоящее время равен нулю, можно пока не учитывать. Мы вернемся к нему позже, когда будем обсуждать кэш kv.

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

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

Эмбеддинг — это фиксированное векторное представление каждого токена, которое больше подходит для глубокого обучения, чем простые целые числа, поскольку оно отражает семантическое значение слов. Размер этого вектора — размерность модели, которая варьируется в зависимости от модели. Например, в LLaMA-7B размерность модели составляет n_embd=4096.

Параметры модели включают в себя матрицу встраивания токенов, которая преобразует токены во встраиваемые объекты. Поскольку размер нашего словаря составляет n_vocab=32000, это 32000 x 4096 матрица, в которой каждая строка содержит вектор встраивания для одного токена:

Матрица встраивания токенов

Каждому токену соответствует встраивание, которое было изучено в процессе обучения и доступно в виде матрицы встраиваний токенов.

Первая часть графа вычислений извлекает из матрицы встраивания токенов соответствующие строки для каждого токена.

// llama.cpp (simplified)
static struct ggml_cgraph * llm_build_llama(/* ... */) {
    // ...

    struct ggml_tensor * inp_tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
    memcpy(
        inp_tokens->data,
        tokens,
        n_tokens * ggml_element_size(inp_tokens));

    inpL = ggml_get_rows(ctx0, model.tok_embeddings, inp_tokens);
}
//

The code first creates a new one-dimensional tensor of integers, called inp_tokens, to hold the numerical tokens. Then, it copies the token values into this tensor’s data pointer. Last, it creates a new GGML_OP_GET_ROWS tensor operation combining the token-embedding matrix model.tok_embeddings with our tokens.

This operation, when later computed, pulls rows from the embeddings matrix as shown in the diagram above to create a new n_tokens x n_embd matrix containing only the embeddings for our tokens in their original order:

Матрица встраивания

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

Трансформатор

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

Внимание к себе

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

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

Входными данными для механизма самовнимания является n_tokens x n_embd матрица эмбеддинга, в которой каждая строка, или вектор, представляет отдельный токен4. Затем каждый из этих векторов преобразуется в три отдельных вектора, называемых «ключом», «запросом» и «значением». Преобразование осуществляется путем умножения вектора эмбеддинга каждого токена на фиксированные матрицы wk, wq и wv, которые являются частью параметров модели:

Векторы ключа, запроса и значения

При умножении вектора встраивания токена на матрицы параметров wk, wq и wv получаются векторы «ключ», «запрос» и «значение» для этого токена.

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

// llama.cpp (simplified to remove use of cache)

// `cur` contains the input to the self-attention mechanism
struct ggml_tensor * K = ggml_mul_mat(ctx0,
    model.layers[il].wk, cur);
struct ggml_tensor * Q = ggml_mul_mat(ctx0,
    model.layers[il].wq, cur);
struct ggml_tensor * V = ggml_mul_mat(ctx0,
    model.layers[il].wv, cur);

В итоге мы получаем K, Q и V: три матрицы размером n_tokens x n_embd, в которых векторы ключей, запросов и значений для каждого токена объединены.

Следующий этап самовнимания заключается в умножении матрицы Q, содержащей сложенные векторы запросов, на транспонированную матрицу K, содержащую сложенные векторы ключей. Для тех, кто не очень хорошо знаком с матричными операциями, поясним, что эта операция по сути вычисляет совместную оценку для каждой пары векторов запросов и ключей. Мы будем использовать обозначение S(i,j) для обозначения оценки запроса i с ключом j.

В результате этого процесса получается n_tokens^2 оценок, по одной для каждой пары «запрос — ключ», которые объединяются в одну матрицу под названием KQ. Эта матрица впоследствии маскируется, чтобы удалить элементы выше диагонали:

Векторы ключа, запроса и значения

Совместный показатель S(i,j) рассчитывается для каждой пары «запрос — ключ» путем умножения Q на транспонированную матрицу K. Здесь показан результат для первых четырех токенов, а также для токенов, представленных каждым показателем. На этапе маскировки сохраняются только показатели между токеном и предыдущими токенами. Для простоты промежуточная операция масштабирования опущена.

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

Последний этап самовнимания заключается в умножении замаскированного балла KQ_masked на векторы значений, полученные ранее5. Такая операция матричного умножения создаёт взвешенную сумму векторов значений всех предыдущих токенов, где весами являются баллы S(i,j). Например, для четвёртого токена ics создаётся взвешенная сумма векторов значений Quant, um, ▁mechan и ics с весами S(3,0)S(3,3), которые сами были рассчитаны на основе вектора запроса ics и всех предыдущих векторов ключей.

Матрица KQV

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

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

Слои трансформатора

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

  • В сети прямого распространения используются большие фиксированные матрицы параметров. В LLaMA-7B их размеры составляют n_embd x n_ff = 4096 x 11008.
  • Помимо самовнимания, все остальные операции можно рассматривать как выполняемые построчно или токен за токеном. Как упоминалось ранее, только самовнимание содержит межтокеновые вычисления. Это будет важно позже, при обсуждении kv-кэша.
  • Входная и выходная данные всегда имеют размер n_tokens x n_embd: по одной строке для каждого токена, размер каждой из которых соответствует размерности модели.

Для полноты картины я приложил схему одного слоя трансформера в LLaMA-7B. Обратите внимание, что в будущих моделях точная архитектура, скорее всего, будет немного отличаться.

Полный слой трансформера

Full computation graph of a Transformer layer in LLaMA-7B, containing self-attention and feed-foward mechanisms. The output of each layer serves as the input to the next. Large parameter matrices are used both in the self-attention stage and in the feed-forward stage. These constitute most of the 7 billion parameters of the model.

В архитектуре Transformer есть несколько слоёв. Например, в LLaMA-7B их n_layers=32 штук. Слои идентичны, за исключением того, что у каждого из них свой набор матриц параметров (например, свои матрицы wk, wq и wv для механизма самовнимания). Входными данными для первого слоя является матрица встраивания, как описано выше. Выходные данные первого слоя используются в качестве входных данных для второго слоя и так далее. Можно представить, что каждый слой создаёт список векторных представлений, но каждое векторное представление больше не привязано напрямую к отдельному токену, а связано с более сложным пониманием взаимосвязей между токенами.

Вычисление логитов

На последнем этапе работы трансформера вычисляются логиты. Логит — это число с плавающей запятой, которое представляет собой вероятность того, что конкретный токен является «правильным» следующим токеном. Чем выше значение логита, тем больше вероятность того, что соответствующий токен является «правильным».

Логиты вычисляются путём умножения выходных данных последнего слоя Transformer на фиксированную n_embd x n_vocab матрицу параметров (также называемую output в llama.cpp). В результате этой операции для каждого токена в нашем словаре получается логит. Например, в LLaMA получается n_vocab=32000 логитов:

Расчет логита

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

Логиты — это выходные данные трансформера, которые сообщают нам, какие токены являются наиболее вероятными следующими. На этом все тензорные вычисления завершаются. Следующая упрощённая и прокомментированная версия функции llm_build_llama() обобщает все шаги, описанные в этом разделе:

// llama.cpp (simplified and commented)

static struct ggml_cgraph * llm_build_llama(
         llama_context & lctx,
     const llama_token * tokens,
                   int   n_tokens,
                   int   n_past) {
    ggml_cgraph * gf = ggml_new_graph(ctx0);
    struct ggml_tensor * cur;
    struct ggml_tensor * inpL;

    // Create a tensor to hold the tokens.
    struct ggml_tensor * inp_tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, N);
    // Copy the tokens into the tensor
    memcpy(
        inp_tokens->data,
        tokens,
        n_tokens * ggml_element_size(inp_tokens));

    // Create the embedding matrix.
    inpL = ggml_get_rows(ctx0,
        model.tok_embeddings,
        inp_tokens);

    // Iteratively apply all layers.
    for (int il = 0; il < n_layer; ++il) {
        struct ggml_tensor * K = ggml_mul_mat(ctx0, model.layers[il].wk, cur);
        struct ggml_tensor * Q = ggml_mul_mat(ctx0, model.layers[il].wq, cur);
        struct ggml_tensor * V = ggml_mul_mat(ctx0, model.layers[il].wv, cur);

        struct ggml_tensor * KQ = ggml_mul_mat(ctx0, K, Q);
        struct ggml_tensor * KQ_scaled = ggml_scale_inplace(ctx0, KQ, KQ_scale);
        struct ggml_tensor * KQ_masked = ggml_diag_mask_inf_inplace(ctx0,
            KQ_scaled, n_past);
        struct ggml_tensor * KQ_soft_max = ggml_soft_max_inplace(ctx0, KQ_masked);
        struct ggml_tensor * KQV = ggml_mul_mat(ctx0, V, KQ_soft_max);

        // Run feed-forward network.
        // Produces `cur`.
        // ...

        // input for next layer
        inpL = cur;
    }

    cur = inpL;

    // Calculate logits from last layer's output.
    cur = ggml_mul_mat(ctx0, model.output, cur);

    // Build and return the computation graph.
    ggml_build_forward_expand(gf, cur);
    return gf;
}

Для фактического выполнения логического вывода вычисляется граф вычислений, возвращаемый этой функцией, с использованием ggml_graph_compute() как описано ранее. Затем логиты копируются из data указателя последнего тензора в массив чисел с плавающей запятой, готовый к следующему шагу — выборке.

Отбор проб

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

Жадная выборка

Жадная выборка — это простой подход, при котором выбирается токен с наибольшим логитом.

Для нашего примера наиболее высокие логиты имеют следующие токены:

жетон

логит

▁describes

18.990

▁provides

17.871

▁explains

17.403

▁de

16.361

▁gives

15.007

Таким образом, при жадной выборке в качестве следующего токена детерминированно выбирается ▁describes Жадная выборка наиболее полезна, когда при повторной оценке идентичных запросов требуются детерминированные результаты.

Отбор проб температуры

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

  1. Логиты сортируются от большего к меньшему и нормализуются с помощью функции softmax, чтобы их сумма равнялась 1. Это преобразование преобразует каждый логит в вероятность.
  2. Применяется пороговое значение (по умолчанию установлено на уровне 0,95), при котором сохраняются только самые популярные токены, совокупная вероятность которых ниже порогового значения. Этот шаг эффективно удаляет токены с низкой вероятностью, предотвращая случайную выборку «плохих» или «неправильных» токенов.
  3. Оставшиеся логиты делятся на параметр температуры и снова нормализуются таким образом, чтобы их сумма равнялась 1 и они представляли собой вероятности.
  4. На основе этих вероятностей случайным образом выбирается токен. Например, в нашем запросе токен ▁describes имеет вероятность p=0.6, то есть он будет выбран примерно в 60 % случаев. При повторной оценке могут быть выбраны другие токены.

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

Измерение температуры с использованием различных значений

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

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

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

Оптимизирующий вывод

Этап самовнимания в трансформере может стать узким местом с точки зрения производительности по мере роста списка входных токенов для языковой модели. Чем длиннее список токенов, тем больше матриц перемножается. Каждое матричное умножение состоит из множества более мелких числовых операций, известных как операции с плавающей запятой, которые ограничены количеством операций с плавающей запятой в секунду (flops) на графическом процессоре. В арифметике логического вывода Transformer подсчитано, что для модели с 52 миллиардами параметров на графическом процессоре A100 производительность начинает снижаться при 208 токенах из-за чрезмерного количества операций с плавающей запятой. Наиболее часто используемый метод оптимизации для устранения этого узкого места известен как kv-кэш.

Кэш KV

Подведём итог: каждому токену соответствует вектор встраивания, который затем преобразуется в векторы ключа и значения путём умножения на матрицы параметров wk и wv. Кэш kv — это кэш для этих векторов ключа и значения. Кэшируя их, мы экономим на операциях с плавающей запятой, необходимых для их повторного вычисления на каждой итерации.

Кэш работает следующим образом:

  • Во время первой итерации для всех токенов вычисляются векторы ключей и значений, как описано выше, а затем сохраняются в кэше kv.
  • На последующих итерациях нужно будет вычислить только векторы ключей и значений для самого нового токена. Кэш-векторы k-v вместе с векторами k-v для нового токена объединяются в матрицы K и V . Это позволяет избежать повторного вычисления векторов k-v для всех предыдущих токенов, что может быть весьма существенно.

Кэш - память KV

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

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

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

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

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

  • Векторы значений предыдущих токенов.
  • Баллы рассчитываются на основе ключевых векторов предыдущих токенов.

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

Зависимости расчета матрицы KQV

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

Дальнейшая оптимизация последующих итераций

Вы можете задаться вопросом, почему мы не кэшируем векторы запросов, ведь мы кэшируем векторы ключей и значений. Дело в том, что на самом деле, кроме вектора запроса текущего токена, векторы запросов предыдущих токенов на последующих итерациях не нужны. При наличии кэша kv мы можем использовать механизм самовнимания только с вектором запроса последнего токена. Этот вектор запроса умножается на кэшированную K матрицу для вычисления совместных оценок последнего токена и всех предыдущих токенов. Затем она умножается на кэшированную V матрицу, чтобы вычислить только последнюю строку матрицы KQV . Фактически на всех уровнях мы теперь передаём векторы размером 1 x n_embd вместо матриц n_token x n_embd , вычисленных на первой итерации. Чтобы проиллюстрировать это, сравните следующую диаграмму, на которой показана более поздняя итерация, с предыдущей:

Матрица KQV на последующих итерациях

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

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

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

Кэш KV на практике

Мы можем изучить код llama.cpp, чтобы увидеть, как на практике реализован кэш kv. Возможно, вас не удивит, что он создан с использованием тензоров: один для векторов ключей, другой для векторов значений.

// llama.cpp (simplified)

struct llama_kv_cache {
    // cache of key vectors
    struct ggml_tensor * k = NULL;

    // cache of value vectors
    struct ggml_tensor * v = NULL;

    int n; // number of tokens currently in the cache
};

При инициализации кэша выделяется достаточно места для хранения 512 векторов ключей и значений для каждого слоя.

// llama.cpp (simplified)
// n_ctx = 512 by default
static bool llama_kv_cache_init(
    struct llama_kv_cache & cache,
    ggml_type   wtype,
    int   n_ctx) {
    // Allocate enough elements to hold n_ctx vectors for each layer.
    const int64_t n_elements = n_embd*n_layer*n_ctx;

    cache.k = ggml_new_tensor_1d(cache.ctx, wtype, n_elements);
    cache.v = ggml_new_tensor_1d(cache.ctx, wtype, n_elements);

    // ...
}

Напомним, что во время вывода график вычислений строится с помощью функции llm_build_llama(). У этой функции есть параметр с именем n_past, который мы игнорировали ранее. На первой итерации n_tokens параметр содержит количество токенов и n_past имеет значение 0. В последующих итерациях, n_tokens устанавливается равным 1, потому что обрабатывается только последний токен, и n_past содержит количество прошлых токенов. n_pastзатем используется для извлечения правильного количества векторов ключей и значений из кэша kv.

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

// llama.cpp (simplified and commented)

static struct ggml_cgraph * llm_build_llama(
    llama_context & lctx,
    const llama_token * tokens,
    int   n_tokens,
    int   n_past) {
    // ...

    // Iteratively apply all layers.
    for (int il = 0; il < n_layer; ++il) {
         // Compute the key vector of the latest token.
         struct ggml_tensor * Kcur = ggml_mul_mat(ctx0, model.layers[il].wk, cur);
         // Build a view of size n_embd into an empty slot in the cache.
         struct ggml_tensor * k = ggml_view_1d(
            ctx0,
            kv_cache.k,
            // size
            n_tokens*n_embd,
            // offset
            (ggml_element_size(kv_cache.k)*n_embd) * (il*n_ctx + n_past)
         );

         // Copy latest token's k vector into the empty cache slot.
         ggml_cpy(ctx0, Kcur, k);

         // Form the K matrix by taking a view of the cache.
         struct ggml_tensor * K =
             ggml_view_2d(ctx0,
                 kv_self.k,
                 // row size
                 n_embd,
                 // number of rows
                 n_past + n_tokens,
                 // stride
                 ggml_element_size(kv_self.k) * n_embd,
                 // cache offset
                 ggml_element_size(kv_self.k) * n_embd * n_ctx * il);
    }
}

Сначала вычисляется новый вектор ключа. Затем n_past используется для поиска следующего свободного слота в кэше, и туда копируется новый вектор ключа. Наконец, матрица K формируется путём просмотра кэша с правильным количеством токенов (n_past + n_tokens).

Кэш kv является основой для оптимизации логического вывода в LLM. Стоит отметить, что версия, реализованная в llama.cpp (на момент написания статьи) и представленная здесь, не является наиболее оптимальной. Например, она заранее выделяет большой объём памяти для хранения максимального количества поддерживаемых векторов ключей и значений (в данном случае 512). Более продвинутые реализации, такие как vLLM, направлены на повышение эффективности использования памяти и могут обеспечить дальнейшее повышение производительности. Эти продвинутые методы будут рассмотрены в следующих публикациях. Более того, поскольку эта сфера развивается с молниеносной скоростью, в будущем, вероятно, появятся новые и более совершенные методы оптимизации.

Заключительный

В этой статье мы рассмотрели довольно много тем и надеемся, что вы получили базовое представление о полном процессе логического вывода в LLM. Обладая этими знаниями, вы сможете работать с более продвинутыми ресурсами:

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

  • Квантованные модели.
  • Тонкая настройка LLM с помощью LoRA.
  • Различные механизмы внимания (многозадачное внимание, внимание к сгруппированным запросам и скользящее внимание).
  • Пакетная обработка запросов LLM.
  • Грамматическая выборка.

Оставайтесь с нами!


Примечания

Примечания

  1. ggml также предоставляет ggml_build_backward() функцию, которая вычисляет градиенты в обратном направлении — от вывода к вводу. Эта функция используется для обратного распространения ошибки только во время обучения модели, но не во время логического вывода.

  2. В статье описывается модель «кодировщик-декодировщик». LLaMA — это модель только с декодером, поскольку она предсказывает только один токен за раз. Но основные концепции те же.

  3. Для простоты я решил описать здесь механизм самовнимания с одной головкой. В LLaMA используется механизм самовнимания с несколькими головками. Он лишь немного усложняет тензорные операции, но не влияет на основные идеи, представленные в этом разделе.

  4. Если быть точным, то сначала векторные представления подвергаются операции нормализации, которая масштабирует их значения. Мы не учитываем этот фактор, поскольку он не влияет на основные представленные идеи.

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment