Оптимізація інференсу LLM через TensorRT-LLM

Оптимізація інференсу LLM за допомогою TensorRT-LLM

TensorRT-LLM — бібліотека NVIDIA для максимально ефективного LLM-інференсу на NVIDIA GPU. Якщо vLLM - зручний production сервер, то TensorRT-LLM - низькорівневий двигун для максимальної продуктивності на обладнанні NVIDIA. Прискорення: 2-4x в порівнянні з vLLM на тих же GPU.

Архітектура та принцип роботи

TensorRT-LLM компілює модель оптимізований TensorRT двигун:

1. Graph compilation: граф моделі компілюється з урахуванням конкретного GPU (архітектура, VRAM, тензорні ядра)
2. Kernel fusion: кілька операцій об'єднуються в один CUDA-kernel (LayerNorm + Linear, Flash Attention)
3. Quantization: FP8, INT8, INT4 з точними calibration методами
4. In-flight batching: найбільш просунута реалізація continuous batching

Встановлення та компіляція моделі

# Установка через Docker (рекомендуется)
docker pull nvcr.io/nvidia/tritonserver:24.05-trtllm-python-py3

# Или pip
pip install tensorrt-llm --extra-index-url https://pypi.nvidia.com

import tensorrt_llm
from tensorrt_llm.builder import BuildConfig, build_model
from tensorrt_llm.models import LLaMAForCausalLM

# Загрузка HuggingFace модели
hf_model_path = "meta-llama/Llama-3-8b-instruct"

# Конфигурация компиляции
build_config = BuildConfig(
    max_batch_size=64,
    max_input_len=2048,
    max_output_len=512,
    max_beam_width=1,               # greedy decoding
    strongly_typed=True,
    plugin_config={
        "gpt_attention_plugin": "float16",
        "gemm_plugin": "float16",
        "rmsnorm_quantization_plugin": False,
        "use_paged_context_fmha": True,    # PagedAttention
        "use_fp8_context_fmha": False,
    }
)

# Компиляция занимает 5-30 минут в зависимости от модели и GPU
engine = build_model(
    model=LLaMAForCausalLM.from_hugging_face(hf_model_path),
    build_config=build_config
)
engine.save("./llama3-8b-engine/")

FP8 Quantization на H100

H100 має апаратну підтримку FP8 – найбільший приріст продуктивності:

from tensorrt_llm.quantization import QuantAlgo

build_config_fp8 = BuildConfig(
    max_batch_size=128,
    max_input_len=4096,
    max_output_len=1024,
    quant_config=QuantConfig(
        quant_algo=QuantAlgo.FP8,
        kv_cache_quant_algo=QuantAlgo.FP8,  # KV-кеш тоже в FP8
    ),
    plugin_config={
        "use_fp8_context_fmha": True,  # Flash Attention в FP8
        "gemm_plugin": "float16",
    }
)

FP8 на H100: приблизно 2x приріст черезвихід у порівнянні з BF16, деградація якості < 0.5% на стандартних бенчмарках.

Інтеграція з Triton Inference Server

TensorRT-LLM нативно інтегрується з NVIDIA Triton:

# Структура для Triton
model_repository/
├── ensemble/
│   └── config.pbtxt
├── preprocessing/      # токенизация
│   ├── config.pbtxt
│   └── 1/model.py
├── tensorrt_llm/       # TRT-LLM движок
│   ├── config.pbtxt
│   └── 1/
│       ├── model.engine
│       └── config.json
└── postprocessing/     # детокенизация
    ├── config.pbtxt
    └── 1/model.py

# tensorrt_llm/config.pbtxt
name: "tensorrt_llm"
backend: "tensorrtllm"
max_batch_size: 128

parameters {
  key: "max_beam_width"
  value: { string_value: "1" }
}
parameters {
  key: "executor_worker_path"
  value: { string_value: "/opt/tritonserver/backends/tensorrtllm/trtllmExecutorWorker" }
}
parameters {
  key: "decoding_mode"
  value: { string_value: "top_p_top_k" }
}

Multi-GPU з Tensor Parallelism

# LLaMA-70B на 4xH100
build_config_tp4 = BuildConfig(
    max_batch_size=64,
    max_input_len=8192,
    max_output_len=2048,
    auto_parallel_config=AutoParallelConfig(
        world_size=4,
        gpus_per_node=4,
        shards_along_head=4,       # tensor parallelism
    )
)

# Запуск mpirun для multi-GPU
# mpirun -n 4 python run_inference.py

Порівняння з vLLM

Рекомендація: vLLM для більшості production use cases. TensorRT-LLM - коли потрібно максимально використовувати NVIDIA GPU (високонавантажені сервіси, cost optimization на cloud GPU).

Терміни впровадження

День 1–3: Установка TRT-LLM, компіляція першої моделі, вимірювання baseline метрик

Тиждень 1–2: Підбір оптимальних параметрів компіляції, квантизація, інтеграція з Triton

Тиждень 3–4: Load testing, monitoring, production deployment

Місяць 2: Оптимізація для конкретних use cases (latency vs throughput), multi-model deployment