Giới thiệu thư viện TensorRT-LLM tối ưu cho mô hình ngôn ngữ và cách cài đặt

TensorRT-LLM là một thư viện của Python cho phép tối ưu các mô hình ngôn ngữ và xây dựng TensorRT Engine để thực hiện Inference đồ thị của mô hình ngôn ngữ.

Một số công nghệ song song mà thư viện này sử dụng để tối ưu:

• Tensor Parallelism

• Pipeline Parallelism

Các thuật toán này sẽ chia nhỏ mô hình thành nhiều phần và thực hiện tối ưu tính toán song song trên các phần đó khiến việc inference trở nên nhanh chóng hơn.

Ngoài ra thư viện này cho phép mô hình chạy ở INT4 hoặc INT8 với bộ tham số (weights) và FP16 với activations. Đây là một số kỹ thuật lượng tử hóa đã được đề cập trong bài viết sau. Có một kỹ thuật tên SmoothQuant cũng được đề cập trong thư viện này.

Cách cài đặt thư viện này trên Ubuntu:

• Tải source code của TensorRT-LLM về máy

# TensorRT-LLM uses git-lfs, which needs to be installed in advance.
apt-get update && apt-get -y install git git-lfs

git clone https://github.com/NVIDIA/TensorRT-LLM.git
cd TensorRT-LLM
git submodule update --init --recursive
git lfs install
git lfs pull

• Build Docker từ Source Code này

docker build . -f docker/Dockerfile.multi -t tensorrt_llm/devel:latest

Thời gian build sẽ khoảng từ 2-3 tiếng

• Chạy docker container

docker run --rm -it --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 --gpus=all --volume ${PWD}:/code/tensorrt_llm --workdir /code/tensorrt_llm tensorrt_llm/devel:latest

Ở trong Container này, bạn có thể bắt đầu sử dụng các mô hình ngôn ngữ khác nhau.

Giao diện của TensorRT-LLM đã được cài đặt.

Cách chạy một mô hình bất kỳ

Ví dụ ở đây mình sẽ hướng dẫn các bạn sử dụng mô hình BLOOM

• Cài đặt thư viện

pip install -r examples/bloom/requirements.txt
git lfs install

• Tải Weights

cd examples/bloom
rm -rf ./bloom/560M
mkdir -p ./bloom/560M && git clone https://huggingface.co/bigscience/bloom-560m ./bloom/560M

• Build Engine

# Single GPU on BLOOM 560M
python build.py --model_dir ./bloom/560M/ \
                --dtype float16 \
                --use_gemm_plugin float16 \
                --use_gpt_attention_plugin float16 \
                --output_dir ./bloom/560M/trt_engines/fp16/1-gpu/

• Tiến hành Inference

Bài toán ở đây là bài toán tóm tắt văn bản.

python summarize.py --test_trt_llm \
                    --hf_model_location ./bloom/560M/ \
                    --data_type fp16 \
                    --engine_dir ./bloom/560M/trt_engines/fp16/1-gpu/