Привет, Хабр! На связи CEO команды Compressa AI. Недавно обнаружил для себя крутой базовый курс по эффективному запуску и инференсу LLM моделей от легенды AI мира — Andrew NG и его платформы DeepLearning. Он полностью на английском языке в формате видео, поэтому я осмелился адаптировать его под формат Хабра на русском языке. Знания должны быть доступны всем и в удобной форме, так ведь?

Многие команды (включая и Compressa AI) начинали LLM проекты с использования облачных API. Но по мере развития все больше разработчиков хотят использовать open-source LLM, чтобы экономить на токенах, снижать latency, запускать fine-tuning на собственных данных и в целом меньше зависеть от внешних моделей.

Из этого курса вы узнаете детали эффективного обслуживания и дообучения open-source LLM, включая методы обработки множества запросов от нескольких пользователей. Используя несколько таких методов одновременно, вы можете улучшить как задержку (latency), так и пропускную способность (throughput). Например, благодаря применению последних open-source технологий в своем продукте, мы добились увеличения пропускной способности до 70x на 1 GPU в сравнении с дефолтными Hugging Face & PyTorch.

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

Что внутри курса?

• Детали генерации текста с помощью LLM по одному токену за раз и реализация KV-Caching;

• Batching для обработки нескольких входных данных одновременно;

• Continuous batching для обработки потока запросов в реальном времени без ожидания формирования полных батчей;

• Quantization для уменьшения потребления памяти моделью и снижения требований к железу;

• LoRA как эффективный метод fine-tuning без изменения изначальных весов модели;

• Объединение нескольких LoRA с continuous batching для одновременного обслуживания десятков дообученных моделей.

От слов к… сути курса

Следуя структуре курса, в первом уроке расскажу и покажу, как итеративно генерировать текст с LLM по одному токену за раз и как этот процесс разделить на две фазы — pre-fill и decode, а также оптимизировать его с помощью KV-caching (кеширования части вычислений).

Загрузка LLM из Hugging Face

Начнем с загрузки LLM из Hugging Face, которая послужит примером модели для инференса (a.k.a. генерации output токенов). В оригинальном курсе используется GPT-2, но в ней плохо поддерживается русский язык, поэтому буду использовать русскоязычный аналог — rugpt3small_based_on_gpt2 на протяжении всего курса.

Давайте теперь импортируем необходимые зависимости (PyTorch, transformers и другие), а затем загрузим нашу LLM и соответствующий токенизатор из Hugging Face.

# Импортируем зависимости
import matplotlib.pyplot as plt
import numpy as np
import time
import torch

# Загружаем LLM и токенизатор
from transformers import AutoModelForCausalLM, AutoTokenizer
model_name = "ai-forever/rugpt3small_based_on_gpt2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

Если прилетают ошибки, прочитайте их и установите нужные библиотеки с помощью pip install или аналога.

А если хотите глубже посмотреть на архитектуру модели, используйте команду ниже:

print(model)

Здесь стоит напомнить, что GPT-2 и большинство современных LLM — это decoder модели с точки зрения архитектуры. Если изначально модели следовали логике «encoder преобразует input токены в embeddings, а decoder генерирует output токены на основе этих embeddings», то в GPT-2 входные данные сразу преобразуются в embeddings, а затем проходят через серию блоков. Ключевая особенность таких моделей — генерация текста по одному токену за раз, что делает их авторегрессионными. Визуализация на картинке ниже:

Процесс генерации токенов

Ок, теперь когда поговорили про особенности нашей модели и загрузили ее, давайте исследуем сам процесс генерации текста. Возьмем простой промпт:

prompt = "Каждое утро я пью кофе с"

Теперь пропустим этот промпт через токенизатор:

inputs = tokenizer(prompt, return_tensors="pt")

Здесь return_tensors="pt" указывает, что мы хотим получить результаты в формате PyTorch. Он отличается от форматов NumPy или TensorFlow, которые вы тоже можете использовать, но сегодня работаем именно с ним.

Результат токенизации — словарь, содержащий несколько тензоров:

print(inputs)
# {'input_ids': tensor([[ 6207,   900,   647,  7949,   417, 40946,  6715,   281]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1]])}

Разберем полученные тензоры:

• input_ids: Основной тензор, представляющий числовое отображение нашего текста в токены. Каждое число соответствует определенному токену в словаре модели.

• attention_mask: Этот тензор состоит из единиц и имеет ту же длину, что и input_ids.

Концепция attention_mask станет более понятной, когда мы перейдем к batching. Пока достаточно воспринимать его как вспомогательный тензор, который идет вместе с input_ids и определяет, на какие токены внутри входных данных LLM следует обращать внимание при обработке.

После токенизации входных данных мы готовы передать их в модель и проанализировать результат. Как это делается:

# Отключаем вычисление градиентов для экономии памяти при инференсе
with torch.no_grad():
    # Передаем токенизированные входные данные в модель
    outputs = model(**inputs)

# Извлекаем logits — необработанные предсказания модели
logits = outputs.logits
print(logits.shape)
# torch.Size([1, 8, 50264])

Полученный тензор имеет 3 измерения:

1: размер batch (в нашем случае 1, так как мы передали один input);
8: количество токенов в нашем input;
50264: размер словаря модели (количество возможных токенов для output).

Следующий шаг после получения logits от модели — определить, какой токен модель предскажет как наиболее вероятное продолжение последовательности:

# Выбираем logits только для нового токена
last_logits = logits[0, -1, :]
# Находим индекс токена с наибольшей вероятностью быть следующим в последовательности
next_token_id = last_logits.argmax()
print(next_token_id)
# tensor(29258)

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

predicted_token = tokenizer.decode(next_token_id)
print(predicted_token)
# ' молоком'

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

Теперь вместо выбора только наиболее вероятного токена давайте рассмотрим топ-10 наиболее вероятных вариантов продолжения:

top_k = torch.topk(last_logits, k=10)
tokens = [tokenizer.decode(tk) for tk in top_k.indices]
print(tokens)
# [' молоком', ' сахаром', ' лим', ' кори', ' шоколад', ' конья', ' було', ' м', ' пирож', ' утра']

Ставьте лайк, если тоже любите кофе с конья(ком) по утрам :)

При использовании разных стратегий декодирования модель может выбрать какой-то из этих альтернативных токенов. В LLM степень «креативности» при выборе менее вероятных токенов обычно контролируется температурой. Но в это не буду углубляться в рамках текущего урока.

Собственно, после генерации первого токена для продолжения нашей фразы (молоком) можем использовать его для создания нового входного тензора и дальнейшей генерации текста:

next_inputs = {
    # Обновляем input_ids, добавляя новый токен к исходной последовательности
    "input_ids": torch.cat(
        [inputs["input_ids"], next_token_id.reshape((1, 1))],
        dim=1
    ),

    # Обновляем attention_mask, добавляя 1 для нового токена

    "attention_mask": torch.cat(
        [inputs["attention_mask"], torch.tensor([[1]])],
        dim=1
    )
}

После добавления нового токена давайте взглянем на обновленный input:

print(next_inputs["input_ids"], next_inputs["input_ids"].shape)
print(next_inputs["attention_mask"], next_inputs["attention_mask"].shape)
# tensor([[ 6207,   900,   647,  7949,   417, 40946,  6715,   281, 29258]]) torch.Size([1, 9])
# tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1]]) torch.Size([1, 9])

Если разобрать результаты:

• Видим новые input_ids, включая новый токен (29258) в конце;

• Форма тензора изменилась с [1, 8] на [1, 9], отражая добавление нового токена;

• Аналогично для attention_mask.

Замеряем скорость

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

Начнем с определения функции, которая объединит все, что мы делали в предыдущих ячейках. Она будет принимать один словарь inputs и генерировать следующий токен:

def generate_token(inputs):
    with torch.no_grad():
        outputs = model(**inputs)

    logits = outputs.logits
    last_logits = logits[0, -1, :]
    next_token_id = last_logits.argmax()
    return next_token_id

Теперь, когда у нас есть вспомогательная функция generate_token, попробуем сгенерировать несколько токенов и посмотрим, сколько времени это займет:

generated_tokens = []  # Список для хранения сгенерированных токенов в виде текста
next_inputs = inputs   # Входные данные, которые будут обновляться на каждом шаге
durations_s = []       # Список для хранения длительности каждой итерации в секундах

# Генерация 20 токенов
for _ in range(20):
    # Замер времени перед генерацией токена
    t0 = time.time()
    
    # Генерация следующего токена
    next_token_id = generate_token(next_inputs)
    
    # Запись времени, затраченного на операцию
    durations_s += [time.time() - t0]
    
    # Обновление входных данных для следующей итерации
    next_inputs = {
        "input_ids": torch.cat(
            [next_inputs["input_ids"], next_token_id.reshape((1, 1))],
            dim=1),  # Конкатенация input_ids с новым token_id
        "attention_mask": torch.cat(
            [next_inputs["attention_mask"], torch.tensor([[1]])],
            dim=1)  # Добавление 1 к attention_mask
    }
    
    # Декодирование token_id в текст и добавление в список
    next_token = tokenizer.decode(next_token_id)
    generated_tokens.append(next_token)

# Вывод общего времени генерации
print(f"{sum(durations_s)} s")
# Вывод списка сгенерированных токенов
print(generated_tokens)
# 0.540771484375 s
# [' молоком', '.', ' ', ' И', ' каждый', ' вечер', ' я', ' пью', ' кофе', ' с', ' молоком', '.', ' ', ' И', ' каждый', ' вечер', ' я', ' пью', ' кофе', ' с']

Результат показывает, что генерация 20 токенов заняла около 0.54 секунд. Чтобы детальнее проанализировать динамику этой скорости от токена к токену, смотрим наглядный график:

plt.plot(durations_s)
plt.show()

 Замечаем несколько моментов:

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

2. Генерация первого токена обычно занимает немного больше времени по ряду причин, но не будем сейчас в это углубляться. После этого время генерации отдельного токена резко падает и начинает постепенно расти. Примерно с 200 миллисекунд до 300+ к концу последовательности.

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

Используем кэш для ускорения вычислений

Для трансформеров одной из наиболее ресурсоемких операций является вычисление внимания (attention). Про механизм attention можно написать даже отдельную статью, поэтому постараюсь объяснить кратко и с небольшим упрощением. Attention использует 3 ключевых компонента:

• Q (Query — запрос): «Что нужно сейчас?»;

• K (Key — ключ): «Ярлыки» для слов во входном тексте;

• V (Value — значение): Фактическая информация о словах.

Когда LLM прямолинейно генерирует длинный текст токен за токеном, для каждого нового слова ей приходится «перечитывать» и «обдумывать» весь предыдущий контекст, то есть заново вычислять эти значения для всей последовательности. Это похоже на то, как если бы вы, продолжая длинное предложение, каждый раз перечитывали его с самого начала.

KV-Caching позволяет избежать таких затрат ресурсов и поделить процесс на 2 фазы:

• При генерации первого токена (фаза pre-fill) модель вычисляет Q, K и V для всего входного текста.

• Для последующих токенов (фаза decode):

  • Вычисляются Q, K и V только для нового токена.

  • K и V для предыдущих токенов берутся из кэша.

  • Новые K и V добавляются в кэш.

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

Итак, ранее уже определяли функцию generate_token, которая принимала на вход только input_ids и attention_mask. Теперь создадим новую функцию generate_token_with_past, которая делает то же самое, но использует и возвращает past_key_values — значения K и V, уже вычисленные для данного входа модели.

def generate_token_with_past(inputs):
    with torch.no_grad():
        outputs = model(**inputs)

    logits = outputs.logits
    last_logits = logits[0, -1, :]
    next_token_id = last_logits.argmax()
    return next_token_id, outputs.past_key_values

Давайте теперь повторим процесс, который делали ранее для измерения времени генерации 20 токенов. Но на этот раз будем использовать KV-caching и передавать только следующий токен в качестве input_ids:

generated_tokens = []
next_inputs = inputs
durations_cached_s = []
for _ in range(20):
    t0 = time.time()
    # Используем новую функцию, которая возвращает past_key_values
    next_token_id, past_key_values = generate_token_with_past(next_inputs)
    durations_cached_s += [time.time() - t0]
    
    next_inputs = {
        # Теперь input_ids содержит только новый токен
        "input_ids": next_token_id.reshape((1, 1)),
        "attention_mask": torch.cat(
            [next_inputs["attention_mask"], torch.tensor([[1]])],
            dim=1),
        # Используем закэшированные past_key_values
        "past_key_values": past_key_values
    }
    
    next_token = tokenizer.decode(next_token_id)
    generated_tokens.append(next_token)

print(f"{sum(durations_cached_s)} s")
print(generated_tokens)
# 0.3520181179046631 s
# [' молоком', '.', ' ', ' И', ' каждый', ' вечер', ' я', ' пью', ' кофе', ' с', ' молоком', '.', ' ', ' И', ' каждый', ' вечер', ' я', ' пью', ' кофе', ' с']

Теперь давайте построим наглядный график, который покажет время генерации токенов с применением KV-caching и без него:

plt.plot(durations_s)
plt.plot(durations_cached_s)
plt.show()

Оранжевая линия, похожая на хоккейную клюшку, показывает новые результаты с применением KV-caching. Как и ожидалось, после генерации первого токена время на генерацию каждого последующего резко падает и остается на низком уровне до самого конца.

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

KV-caching оптимизация является первым важным шагом в ускорении инференса LLM. И она лежит в основе того, что делают современные библиотеки. Существуют и другие, более сложные техники оптимизации, которые пытаются эффективно использовать кэш как в памяти, так и в вычислениях на уровне CUDA. В нашей готовой сборке Compressa LLM как раз используются лучшие современные подходы, такие как Page Attention, что позволяет ускорить генерацию для 1 запроса в 2-10 раз.

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

Так когда следующая часть?

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