Оптимизация инференса LLM через Text Generation Inference (TGI)

Оптимизация инференса LLM с помощью Text Generation Inference (TGI)

Text Generation Inference (TGI) — HuggingFace производственный сервер для LLM инференса. Написан на Rust (сервер) + Python (модельная логика). Проще vLLM в конфигурации, нативно интегрирован с HuggingFace Hub. Используется в Hugging Chat и многих production deployment.

Быстрый старт

# Docker (рекомендуется)
docker run --gpus all \
  -p 8080:80 \
  -v /data/models:/data \
  ghcr.io/huggingface/text-generation-inference:2.1 \
  --model-id meta-llama/Llama-3-8b-instruct \
  --max-input-length 4096 \
  --max-total-tokens 8192 \
  --max-batch-prefill-tokens 32768 \
  --num-shard 1 \
  --dtype bfloat16 \
  --huggingface-hub-token $HF_TOKEN

# Клиент через официальный пакет
from huggingface_hub import InferenceClient

client = InferenceClient(model="http://localhost:8080")

response = client.text_generation(
    prompt="Explain transformer attention in simple terms",
    max_new_tokens=512,
    temperature=0.7,
    repetition_penalty=1.1,
    stream=False
)

# Streaming
for token in client.text_generation(prompt, stream=True):
    print(token, end="", flush=True)

Ключевые возможности TGI

Continuous batching (in-flight batching): новые запросы добавляются в батч во время генерации предыдущих. Реализация аналогична vLLM.

Flash Attention 2: эффективная реализация self-attention с O(n) памятью вместо O(n²). Автоматически включается для поддерживаемых моделей.

Tensor Parallelism: распределение модели на несколько GPU через --num-shard.

Speculative Decoding: через --speculate N — draft модель генерирует N токенов, target верифицирует.

Quantization: поддержка GPTQ, AWQ, EETQ, BitsAndBytes из коробки.

Конфигурация для разных сценариев

# Максимальный throughput (batch processing)
docker run --gpus all ghcr.io/huggingface/text-generation-inference:2.1 \
  --model-id mistralai/Mixtral-8x7B-Instruct-v0.1 \
  --num-shard 2 \                           # 2 GPU для MoE модели
  --max-input-length 8192 \
  --max-total-tokens 16384 \
  --max-batch-prefill-tokens 131072 \       # большой prefill batch
  --max-waiting-tokens 20 \                 # ждём больше токенов для батча
  --dtype bfloat16

# Минимальная latency (интерактивные чаты)
docker run --gpus all ghcr.io/huggingface/text-generation-inference:2.1 \
  --model-id meta-llama/Llama-3-8b-instruct \
  --max-input-length 2048 \
  --max-total-tokens 4096 \
  --max-batch-prefill-tokens 4096 \
  --max-concurrent-requests 32 \            # ограничиваем очередь
  --waiting-served-ratio 1.2

# Экономия VRAM через quantization
docker run --gpus all ghcr.io/huggingface/text-generation-inference:2.1 \
  --model-id TheBloke/Llama-2-13B-AWQ \
  --quantize awq \
  --dtype float16

Custom Handlers

TGI позволяет добавить preprocessing/postprocessing через custom handler:

# custom_handler.py
class CustomHandler:
    def __init__(self):
        self.tokenizer = AutoTokenizer.from_pretrained(...)

    def preprocess(self, inputs: dict) -> dict:
        """Преобразование входящего запроса перед inference."""
        prompt = inputs.get("inputs", "")

        # Добавление system prompt
        full_prompt = f"<|system|>You are a helpful assistant.<|end|>\n<|user|>{prompt}<|end|>\n<|assistant|>"

        return {"inputs": full_prompt, **{k: v for k, v in inputs.items() if k != "inputs"}}

    def postprocess(self, model_output: dict) -> dict:
        """Постобработка вывода модели."""
        generated = model_output["generated_text"]
        # Убираем системный prompt из вывода
        return {"generated_text": generated.split("<|assistant|>")[-1].strip()}

Мониторинг и метрики

TGI экспортирует Prometheus метрики на /metrics:

tgi_request_duration_seconds_bucket  # latency histogram
tgi_batch_inference_duration_seconds  # batch inference time
tgi_request_input_length              # длины входов
tgi_request_generated_tokens          # длины сгенерированных токенов
tgi_batch_current_size                # текущий размер батча
tgi_queue_size                        # размер очереди ожидания

TGI vs vLLM

Для большинства use cases оба варианта дают близкую производительность. TGI удобнее при работе в HF экосистеме.