Posted inAI

Tăng tốc Llama-3 với TensorRT-LLM trên Ubuntu: Từ cài đặt đến Production

Artificial Intelligence tutorial - IT technology blog
Artificial Intelligence tutorial - IT technology blog

Xử lý bài toán Latency trong thực tế

Hệ thống chatbot Llama-3 của team mình từng rơi vào tình trạng quá tải vào giờ cao điểm. Dù chạy vLLM trên cụm A100, tốc độ xử lý (TPS) vẫn tụt dốc không phanh khi lượng người dùng tăng. Latency chạm ngưỡng 10 giây cho mỗi phản hồi, khiến trải nghiệm người dùng cực kỳ tệ. Sau nhiều giờ rà soát, mình nhận ra nút thắt cổ chai nằm ở việc chưa khai thác hết tập lệnh chuyên biệt trên GPU NVIDIA.

Để đưa các mô hình ngôn ngữ lớn (LLM) vào vận hành thực tế, việc chỉ cài đặt và chạy (plug-and-play) là chưa đủ. Bạn cần can thiệp sâu vào tầng biên dịch của GPU. Đó là lúc TensorRT-LLM phát huy tác dụng. Nếu bạn muốn tăng tốc độ phản hồi từ 5-10 tokens/giây lên hơn 50 tokens/giây, đây là giải pháp tối ưu nhất hiện nay.

Tại sao TensorRT-LLM lại nhanh vượt trội?

TensorRT-LLM không đơn thuần là một thư viện nạp mô hình như Transformers. Nó đóng vai trò là một compiler, biên dịch toàn bộ kiến trúc mạng neural thành một engine được tối ưu riêng cho từng dòng chip đồ họa. Thay vì giải thích lý thuyết suông, hãy nhìn vào 3 kỹ thuật then chốt giúp nó đạt hiệu suất cao:

  • In-flight Batching: Kỹ thuật này cho phép đẩy request mới vào xử lý ngay khi có một request cũ vừa hoàn thành, thay vì đợi cả batch chạy xong.
  • Quantization (FP8/INT8): Giảm kích thước mô hình để chạy nhanh hơn trên card H100 hoặc A100. Việc này giúp tiết kiệm tới 40% VRAM mà độ chính xác gần như không đổi.
  • Kernel Fusion: Gộp hàng trăm phép toán ma trận nhỏ lẻ thành một khối duy nhất. Điều này giúp giảm đáng kể thời gian truy xuất bộ nhớ (VRAM throughput).

Thiết lập môi trường chuẩn

Kinh nghiệm xương máu là đừng bao giờ cài đặt trực tiếp lên máy host vì TensorRT-LLM rất nhạy cảm với phiên bản thư viện. Chỉ cần lệch một bản vá CUDA là hệ thống sẽ báo lỗi ngay. Docker là lựa chọn an toàn nhất. Dưới đây là cấu hình mình đã test ổn định trên môi trường production:

  • Hệ điều hành: Ubuntu 22.04 LTS
  • Phần cứng: NVIDIA Ampere trở lên (RTX 3090, 4090, A10, A100…)
  • Driver: Phiên bản 535 hoặc mới hơn

Đầu tiên, hãy cài đặt NVIDIA Container Toolkit để Docker có thể giao tiếp với GPU:

curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
    sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
    sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
sudo apt-get update && sudo apt-get install -y nvidia-container-toolkit
sudo systemctl restart docker

Quy trình triển khai 3 bước

1. Chuẩn bị Checkpoint

Bạn không thể nạp trực tiếp file .safetensors từ HuggingFace vào engine. Chúng ta cần chuyển đổi chúng sang định dạng trung gian của TensorRT-LLM. Sử dụng Docker image chính thức từ NVIDIA (NGC) để đảm bảo tính tương thích:

# Lấy mã nguồn và vào container
git clone https://github.com/NVIDIA/TensorRT-LLM.git
cd TensorRT-LLM

# Chuyển đổi trọng số Llama-3-8B
python3 examples/llama/convert_checkpoint.py \
    --model_dir ./llama-3-8b-hf \
    --output_dir ./llama-3-8b-ckpt \
    --dtype float16

2. Build Engine tối ưu

Đây là giai đoạn quan trọng nhất. Lệnh trtllm-build sẽ phân tích phần cứng hiện tại để tạo ra file thực thi tối ưu nhất. Lưu ý thông số max_batch_size; nếu bạn đặt quá cao, VRAM sẽ bị chiếm dụng hết ngay cả khi không có request.

trtllm-build --checkpoint_dir ./llama-3-8b-ckpt \
    --output_dir ./llama-3-8b-engine \
    --gemm_plugin float16 \
    --max_batch_size 4 \
    --max_input_len 2048

3. Vận hành với Triton Inference Server

Để phục vụ API ổn định cho hàng ngàn người dùng, mình luôn dùng Triton Server thay vì chạy script Python thủ công. Triton quản lý hàng đợi (queue) và phân phối tài nguyên GPU rất hiệu quả. Bạn cần tổ chức thư mục theo cấu hình model_repository của NVIDIA để server nhận diện được engine.

docker run --gpus all --rm -p 8000:8000 \
    -v $(pwd)/model_repository:/models \
    nvcr.io/nvidia/tritonserver:24.03-trtllm-py3 \
    tritonserver --model-repository=/models

Lưu ý để tránh lỗi Out of Memory (OOM)

Trong quá trình chạy thực tế, lỗi OOM thường xảy ra khi KV Cache chiếm quá nhiều bộ nhớ. Nếu card của bạn chỉ có 24GB VRAM như RTX 3090, hãy cân nhắc sử dụng Weight-Only Quantization. Kỹ thuật này giúp nén mô hình Llama-3-8B từ 15GB xuống còn khoảng 8GB. Nhờ đó, bạn có thêm không gian để tăng max_input_len cho các đoạn hội thoại dài.

Ngoài ra, hãy luôn kiểm tra hiệu năng bằng công cụ genai-perf. Trong bài test của team mình, việc bật gemm_plugin giúp tăng tốc độ tính toán ma trận thêm 25% so với cài đặt mặc định. Đừng bỏ qua các flag tối ưu nhỏ vì tổng hợp lại chúng sẽ tạo ra sự khác biệt rất lớn về throughput.

Lời kết

Triển khai TensorRT-LLM đòi hỏi sự tỉ mỉ trong việc thiết lập môi trường và cấu hình thông số. Tuy nhiên, kết quả nhận lại là một hệ thống có khả năng chịu tải tốt và phản hồi cực nhanh. Nếu bạn đang xây dựng sản phẩm AI hướng tới người dùng cuối, tối ưu hóa ở tầng sâu như thế này là bước đi bắt buộc để tiết kiệm chi phí hạ tầng và nâng cao chất lượng dịch vụ.

Share:
Tags: