Оптимізація інференсу 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 екосистемі.