Posted inAI

Triển khai Text Generation Inference (TGI) với Docker để phục vụ LLM hiệu năng cao

AI tutorial - IT technology blog
AI tutorial - IT technology blog

Nghẽn cổ chai trong môi trường Production: Khi Inference đơn giản thất bại

Triển khai một mô hình ngôn ngữ lớn (LLM) như Llama 3 cho một thử nghiệm cá nhân thì rất dễ dàng. Bạn chỉ cần tải mô hình bằng thư viện Transformers, bọc nó trong một endpoint FastAPI nhanh gọn, và nó sẽ hoạt động. Tuy nhiên, đội ngũ của tôi đã học được một bài học đắt giá rằng thiết lập này không thể mở rộng (scale). Sáu tháng trước, chúng tôi đã ra mắt một công cụ nội bộ dựa trên RAG, và nó đã sụp đổ chỉ với 50 người dùng đồng thời.

Các triệu chứng rất rõ rệt. Latency tăng vọt lên hơn 30 giây. Lỗi timeout trở thành chuyện thường ngày. Người dùng phải nhìn chằm chằm vào màn hình trống rỗng trong khi server vật lộn để xử lý một yêu cầu duy nhất, rồi sau đó mới đổ ra một khối văn bản khổng lồ cùng một lúc. Chúng tôi nhận ra rằng việc inference dựa trên Python tiêu chuẩn xử lý lưu lượng truy cập production rất kém. Nó thiếu đi sự hiệu quả cần thiết cho một trải nghiệm người dùng mượt mà.

Thủ phạm: Tại sao cách phục vụ truyền thống bị nghẽn

Để khắc phục tình trạng giật lag, chúng tôi phải suy nghĩ lại về cách LLM xử lý dữ liệu. Các server truyền thống sử dụng xử lý tuần tự hoặc static batching. Trong một static batch, server đợi một số lượng yêu cầu nhất định, nhóm chúng lại và xử lý như một thao tác ma trận duy nhất. Điều này không hiệu quả. Nếu một người dùng yêu cầu một bài thơ haiku và một người khác yêu cầu một bài luận 500 từ, yêu cầu ngắn sẽ bị “bắt làm con tin” bởi yêu cầu dài. Điều này lãng phí các chu kỳ GPU đắt đỏ.

Các triển khai Python tiêu chuẩn cũng gặp khó khăn với Global Interpreter Lock (GIL). Sự quá tải này ngăn cản GPU đạt đến mức bão hòa tối đa. Nếu không có Token Streaming, người dùng không thể thấy đầu ra của mô hình khi nó đang được tạo ra, khiến thời gian chờ đợi cảm giác dài hơn nhiều. Chúng tôi cần Continuous Batching. Kỹ thuật này đưa các yêu cầu mới vào batch ngay khi một token được tạo ra cho một yêu cầu hiện có, giúp GPU luôn bận rộn liên tục.

Chọn Stack phù hợp: Tại sao lại là TGI?

Chúng tôi đã đánh giá một vài engine chuyên dụng trước khi chọn stack cho production:

  • FastAPI + Transformers: Dễ xây dựng nhưng thiếu các tối ưu hóa cho số lượng truy cập đồng thời cao như PagedAttention.
  • vLLM: Cực kỳ nhanh và phổ biến nhờ PagedAttention, nhưng tại thời điểm đó, nó có vẻ ít tích hợp chặt chẽ với hệ sinh thái Hugging Face cho các mô hình ngách cụ thể.
  • Text Generation Inference (TGI): Hugging Face xây dựng công cụ này dành riêng cho API production của họ. Nó là một cỗ máy mạnh mẽ được viết bằng Rust, C++ và Python.

TGI cung cấp hỗ trợ gốc cho Flash Attention, PagedAttention và các phương pháp quantization như AWQ và bitsandbytes. Bằng cách chuyển sang TGI, chúng tôi đã cắt giảm 40% chi phí phần cứng và tăng gấp đôi tổng throughput. Nó đã xử lý hơn 200 yêu cầu đồng thời trên cùng một phần cứng mà trước đó đã thất bại ở mức 50.

Triển khai TGI với Docker: Hướng dẫn từng bước

Docker là cách đáng tin cậy nhất để triển khai TGI. Nó đóng gói các driver NVIDIA phức tạp, CUDA 12.1 kernel và các dependency của Rust vào một container di động duy nhất. Điều này loại bỏ cơn đau đầu mang tên “nó chạy tốt trên máy của tôi”.

1. Yêu cầu phần cứng

Bạn cần một GPU NVIDIA với đủ VRAM. Đối với mô hình Llama 3 (8B), hãy nhắm tới ít nhất 16GB VRAM (như RTX 4090 hoặc A10G). Đảm bảo đã cài đặt NVIDIA Container Toolkit để Docker có thể giao tiếp với phần cứng của bạn.

# Kiểm tra xem Docker có nhận diện được GPU không
docker run --rm --gpus all nvidia/cuda:12.1.0-base-ubuntu22.04 nvidia-smi

2. Khởi chạy TGI Container

Chúng ta sẽ sử dụng image chính thức của Hugging Face để chạy Llama-3-8B-Instruct. Nếu bạn đang sử dụng các mô hình bị giới hạn (gated models), hãy chuẩn bị sẵn Hugging Face Hub token của bạn.

model="meta-llama/Meta-Llama-3-8B-Instruct"
volume=$PWD/data
token="hf_token_cua_ban_tai_day"

docker run --gpus all --shm-size 1g -p 8080:80 \
    -v $volume:/data \
    -e HUGGING_FACE_HUB_TOKEN=$token \
    ghcr.io/huggingface/text-generation-inference:2.0 \
    --model-id $model \
    --max-batch-prefill-tokens 2048 \
    --max-total-tokens 4096

Giải thích các thiết lập:

  • --shm-size 1g: Cấp phát bộ nhớ chia sẻ để giao tiếp GPU-to-GPU nhanh chóng.
  • --max-total-tokens: Đặt giới hạn cứng cho tổng độ dài đầu vào cộng với đầu ra.
  • --max-batch-prefill-tokens: Ngăn chặn lỗi Out-of-Memory (OOM) bằng cách giới hạn số lượng token được xử lý trong giai đoạn prompt ban đầu.

3. Kích hoạt Token Streaming

TGI tỏa sáng khi sử dụng Server-Sent Events (SSE). Điều này cho phép giao diện người dùng của bạn hiển thị văn bản theo từng ký tự. Đây là một đoạn mã Python để nhận stream:

import requests
import json

def stream_llm_response(prompt):
    url = "http://localhost:8080/generate_stream"
    data = {
        "inputs": prompt,
        "parameters": {"max_new_tokens": 500, "temperature": 0.7}
    }

    response = requests.post(url, json=data, stream=True)
    
    # Lặp qua từng dòng trong phản hồi
    for line in response.iter_lines():
        if line:
            decoded = line.decode('utf-8')
            if decoded.startswith("data:"):
                json_data = json.loads(decoded[5:])
                print(json_data['token']['text'], end="", flush=True)

stream_llm_response("Continuous Batching là gì?")

Tối ưu hóa cho Production: Bài học thực tế

Vận hành TGI trong một cluster suốt sáu tháng đã dạy chúng tôi một vài thủ thuật tối ưu hóa quan trọng. Nếu hiệu suất của bạn bị chững lại, hãy xem xét ba khu vực này.

Thu nhỏ mô hình bằng Quantization

Khi VRAM bị hạn chế, hãy sử dụng quantization. Thêm --quantize bitsandbytes-nf4 vào lệnh của bạn có thể giảm đáng kể mức sử dụng bộ nhớ. Đối với mô hình 8B, điều này có thể giảm dung lượng VRAM từ khoảng 15GB xuống dưới 6GB. Điều này cho phép bạn chạy các mô hình lớn hơn trên phần cứng rẻ hơn mà hầu như không mất đi tính logic. Nếu bạn muốn tìm hiểu sâu hơn về kỹ thuật giảm kích thước mô hình, hãy xem hướng dẫn về fine-tuning Llama 3 với QLoRA trên GPU dân dụng.

Mở rộng với Tensor Parallelism

Các mô hình lớn hơn 13B tham số hiếm khi vừa với một GPU phổ thông. TGI giúp việc thiết lập đa GPU trở nên đơn giản. Sử dụng cờ --num-shard (ví dụ: --num-shard 2) để chia mô hình trên hai GPU. TGI sẽ tự động xử lý các phép toán phức tạp để phân phối khối lượng công việc.

Theo dõi các chỉ số (Metrics)

TGI có một endpoint /metrics tích hợp sẵn cho Prometheus. Hãy theo dõi sát sao tgi_request_queue_size. Nếu hàng đợi luôn ở mức cao, GPU của bạn đã bị bão hòa. Đây là tín hiệu để bạn khởi chạy thêm một instance TGI khác đằng sau một load balancer. Để xây dựng hệ thống giám sát mô hình ML toàn diện trên production, bạn cần kết hợp nhiều chỉ số hơn là chỉ queue size.

Xây dựng hạ tầng AI tốt hơn

Chuyển từ một script Python sang một engine chuyên dụng như TGI là một bước tiến lớn đối với bất kỳ kỹ sư AI nào. Docker giúp hạ tầng này có thể tái lập và dễ dàng mở rộng. Bằng cách kết hợp tốc độ của Rust với Continuous Batching, TGI cung cấp một cách mạnh mẽ để phục vụ các mô hình mã nguồn mở ở quy mô lớn. Nếu bạn cần quản lý nhiều model backend cùng lúc, hãy cân nhắc xây dựng AI Gateway với LiteLLM để thống nhất các provider sau một interface duy nhất. Hãy tập trung vào việc tinh chỉnh max-batch-size và theo dõi mức độ sử dụng GPU để tận dụng tối đa phần cứng của bạn.

Share:
Tags: