Как работает continuous batching в TGI (Hugging Face Text Generation Inference)?

Краткий тезис

batching|Continuous batching (batching|непрерывный батчинг) — техника оптимизации инференса LLM, при которой запросы обрабатываются не статическими батчами, а на уровне отдельных токенов: в каждой итерации генерации scheduler выбирает, какие запросы продолжать, а какие завершить, и добавляет новые. TGI (Hugging Face Text Generation Inference) реализует эту технику через scheduler|токен-уровневый scheduler с очередью FIFO + priority, что позволяет эффективно утилизировать GPU, но уступает vLLM в управлении памятью из-за отсутствия PagedAttention.

1. Термин: Continuous Batching (непрерывный батчинг)

batching|Continuous batching — это метод обработки запросов к LLM, при котором батч динамически изменяется на каждом шаге генерации. В отличие от static batching (статический батчинг), где все запросы собираются в фиксированный батч и обрабатываются синхронно до завершения самого длинного, continuous batching позволяет:

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

Термин «токен-уровневый scheduler» (token-level scheduler) — компонент, который на каждой итерации решает, какие запросы будут участвовать в forward pass, основываясь на их текущем состоянии (активен, завершён, ожидает).

2. Зачем нужен continuous batching?

Проблема static batching:

• Все запросы в батче должны генерировать одинаковое количество токенов (или ждать самого длинного).
• Если один запрос генерирует 10 токенов, а другой — 1000, GPU простаивает, пока не завершится длинный запрос.
• Невозможно добавить новый запрос, пока не освободится место в батче.

Continuous batching решает эти проблемы:

• GPU постоянно занят полезной работой.
• Увеличивается throughput (пропускная способность) системы.
• Снижается latency (задержка) для коротких запросов.

3. Как работает continuous batching в TGI: токен-уровневый scheduler

TGI использует iteration-level scheduling (планирование на уровне итераций). Каждая итерация генерации (один forward pass модели) scheduler:

1. Получает список всех активных запросов (тех, которые ещё не сгенерировали end-of-sequence токен).
2. Выбирает, какие из них будут участвовать в текущей итерации (обычно все активные, но может быть ограничение по памяти).
3. Для каждого выбранного запроса подаёт на вход модели предыдущий сгенерированный токен (или промпт на первом шаге).
4. Модель возвращает логиты, из которых выбирается следующий токен (обычно через sampling или greedy decoding).
5. Если запрос сгенерировал EOS или достиг максимальной длины, он помечается как завершённый и удаляется из активного набора.
6. На освободившееся место (в пределах batch size) добавляются новые запросы из очереди ожидания.

Этот процесс повторяется, пока есть активные запросы или запросы в очереди.

4. Детали: очередь запросов (FIFO + priority) и роутинг

TGI организует входящие запросы в очередь с политикой FIFO (first-in, first-out) и возможностью приоритезации.

• FIFO — запросы обрабатываются в порядке поступления.
• Priority — можно задать приоритет (например, для платных пользователей), и scheduler будет выбирать запросы с более высоким приоритетом раньше.

Router (роутер) — компонент TGI, который может направлять запросы на разные модели (если запущено несколько инстансов). Это позволяет балансировать нагрузку и использовать разные модели для разных задач.

5. Управление памятью в TGI

TGI использует static memory allocation (статическое выделение памяти) для KV cache (кэш ключей и значений). Для каждого запроса резервируется фиксированный объём памяти под максимальную длину генерации (max_new_tokens). Это приводит к:

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

В отличие от vLLM, которая использует PagedAttention (страничное управление памятью) и динамически выделяет страницы под KV cache, TGI менее эффективен по памяти, что может снижать throughput при большом количестве запросов.

6. Сравнение TGI и vLLM

7. Преимущества и недостатки continuous batching в TGI

Преимущества

• Высокая утилизация GPU при разнородной длине запросов.
• Низкая задержка для коротких запросов (не ждут длинные).
• Простая интеграция с экосистемой Hugging Face (модели, токенизаторы).
• Встроенная поддержка роутинга и приоритетов.

Недостатки

• Менее эффективное управление памятью по сравнению с vLLM.
• При большом batch size может не хватать памяти из-за статического резервирования.
• Отсутствие advanced оптимизаций вроде prefix caching (кэширование префиксов) или speculative decoding (спекулятивная декодировка) — хотя некоторые добавляются.

8. Пример конфигурации TGI для continuous batching

TGI запускается как Docker-контейнер. Основные параметры, влияющие на continuous batching:

docker run --gpus all -p 8080:80 \
  -e MODEL_ID=meta-llama/Llama-2-7b-chat-hf \
  -e MAX_BATCH_PREFILL_TOKENS=4096 \
  -e MAX_BATCH_TOTAL_TOKENS=8192 \
  -e MAX_INPUT_LENGTH=2048 \
  -e MAX_TOTAL_TOKENS=4096 \
  ghcr.io/huggingface/text-generation-inference:latest

• MAX_BATCH_PREFILL_TOKENS — максимальное количество токенов на этапе prefill (обработка промпта) для всего батча.
• MAX_BATCH_TOTAL_TOKENS — максимальное количество токенов (prefill + decode) для всего батча.
• MAX_INPUT_LENGTH — максимальная длина входного промпта.
• MAX_TOTAL_TOKENS — максимальная длина промпта + генерации для одного запроса.

Эти параметры ограничивают размер батча и количество одновременно обрабатываемых запросов, влияя на эффективность continuous batching.

9. Псевдокод работы scheduler в TGI

# Упрощённая логика token-level scheduler
class ContinuousBatchingScheduler:
    def __init__(self, max_batch_size, max_total_tokens):
        self.active_requests = []  # список активных запросов
        self.waiting_queue = deque()  # очередь ожидания (FIFO)
        self.max_batch_size = max_batch_size
        self.max_total_tokens = max_total_tokens

    def add_request(self, request):
        if len(self.active_requests) < self.max_batch_size:
            self.active_requests.append(request)
        else:
            self.waiting_queue.append(request)

    def step(self):
        # 1. Определить, какие запросы продолжать
        to_remove = []
        for req in self.active_requests:
            # Выполнить forward pass для текущего токена
            next_token = model.forward(req.current_token)
            req.generated_tokens.append(next_token)
            if next_token == EOS or len(req.generated_tokens) >= req.max_new_tokens:
                to_remove.append(req)
        
        # 2. Удалить завершённые запросы
        for req in to_remove:
            self.active_requests.remove(req)
            # Освободившееся место — можно добавить новые
            if self.waiting_queue:
                new_req = self.waiting_queue.popleft()
                self.active_requests.append(new_req)
        
        # 3. Проверить лимиты по токенам (prefill + decode)
        total_tokens = sum(req.num_tokens for req in self.active_requests)
        if total_tokens > self.max_total_tokens:
            # Приостановить или отклонить запросы (упрощённо)
            pass

Реальный код TGI написан на Rust с использованием Candle или PyTorch, но логика аналогична.

Пет-проект для закрепления

Задача Сравнить throughput TGI и vLLM на синтетическом наборе запросов разной длины.

Инструменты

• Docker (TGI, vLLM)
• Python (aiohttp для асинхронных запросов)
• Бенчмарк-скрипт (например, на основе lm-evaluation-harness или самописный)

Шаги:

1. Развернуть TGI и vLLM с одинаковой моделью (например, TinyLlama/TinyLlama-1.1B-Chat-v1.0).
2. Сгенерировать 1000 запросов с разной длиной промпта (от 50 до 500 токенов) и разной длиной генерации (от 10 до 200 токенов).
3. Отправить запросы асинхронно, замерить общее время выполнения и количество успешных ответов.
4. Построить графики зависимости throughput (запросов/сек) от batch size и длины запросов.
5. Сравнить использование GPU памяти (через nvidia-smi).

Ожидаемый результат Вы увидите, что vLLM показывает более высокий throughput при большом количестве запросов и длинных генерациях, а TGI может быть быстрее на малых batch size и коротких запросах.

Связь с другими вопросами

Навигация

• Предыдущий: 452
• Следующий: 454
• Индекс: 00. Индекс разборов