1.TensorRT-LLM 如何提升 LLM 模型推理效率
大型語言模型(Large language models,LLM)是基於大量數據進行預訓練的超大型深度學習模型。底層轉換器是一組神經網絡,這些神經網絡由具有 self-attention 的編碼器和解碼器組成。編碼器和解碼器從一系列文本中提取含義,並理解其中的單詞和短語之間的關係。
當前 LLM 模型推理的主要瓶頸是 GPU 顯存資源不足。因此,各類加速框架主要集中於降低 GPU 顯存峯值和提高 GPU 使用率兩大目標。
TensorRT-LLM[1]是 NVIDIA 推出的大語言模型(LLM)推理優化框架。它提供了一組 Python API 用於定義 LLMs,並且使用最新的優化技術將 LLM 模型轉換爲 TensorRT Engines,推理時直接使用優化後的 TensorRT Engines。
TensorRT-LLM 主要利用以下四項優化技術提升 LLM 模型推理效率。
1.1 量化
模型量化技術是通過降低原始模型的精度來減少模型推理時的 GPU 顯存使用。
TensorRT 支持多種模型的多種精度,以下列舉了部分主流模型支持的量化精度。
W8A8 SQ 使用了 SmoothQuant 技術[2],在不降低模型推理準確率的前提下,將模型權重和激活層都降低爲 INT8 精度,顯著減少了 GPU 顯存消耗。
W4A16/W8A16 是指模型權重爲 INT4 或者 INT8,激活層爲 FP16 精度。
W4A16 AWQ 以及 W4A16 GPTQ 分別實現了 AWQ[3]和 GPTQ[4]兩篇論文中提到的量化方法。模型權重爲 INT4,激活層爲 FP16 精度。
1.2 In-Flight Batching
傳統的 Batching 技術爲 Static Batching 的,需要等 Batching 中所有序列推理完成後才能進行下一次批次。下圖爲一個輸出最大 Token 爲 8,Batch size 爲 4 的推理過程,使用 Static Batching 技術。S3 序列在 T5 時刻就已經完成推理,但是需要等到 S2 序列在 T8 時刻推理完成後纔會處理下一個 sequence,存在明顯的資源浪費。
In-Flight Batching 又名 Continuous Batching 或 iteration-level batching,該技術可以提升推理吞吐率,降低推理時延。Continuous Batching 處理過程如下,當 S3 序列處理完成後插入一個新序列 S5 進行處理,提升資源利用率。詳情可參考論文 Orca: A Distributed Serving System for Transformer-Based Generative Models[5]。
1.3 Attention
Attention 機制用於從序列中提取關鍵/重要信息,在情感識別、翻譯、問答等任務中起着至關重要的作用。Attention 機制按照演進順序可以分爲 MHA(Multi-head Attention)、MQA(Multi-query Attention)[6]以及 GQA(Group-query Attention)[7]機制。MQA 和 GQA 都是 MHA 的變種。
MHA 是標準的多頭注意力機制,每個 query 存儲一份 KV,因此需要使用較多的顯存。MQA 所有 query 共享一份 KV,推理時容易丟失一些細節信息。GQA 將 query 進行分組,組內共享一份 KV,可以有效避免 MHA 和 MQA 的問題。
TensorRT-LLM 支持 MHA、MQA 及 GQA 方式,可以在 tensorrt_llm.functional.gpt_attention 查看具體實現。
1.4 Graph Rewriting
TensorRT-LLM 在將 LLM 模型編譯爲 TensorRT Engines 時會對神經網絡進行優化,提升執行效率。
2.基於阿里雲容器服務 ACK 的實戰體驗
2.1 雲原生 AI 套件
雲原生 AI 套件是阿里雲容器服務 ACK 提供的雲原生 AI 技術和產品方案,幫助企業更快、更高效地落地雲原生 AI 系統。
本文將介紹如何基於阿里雲容器服務 ACK 雲原生 AI 套件,利用 TensorRT-LLM 優化 LLM 模型推理。
2.2 環境配置
1. 參考文檔安裝雲原生 AI 套件[8]。
2. 登陸容器服務管理控制檯[9],在左側導航欄選擇集羣 > 應用 > 雲原生 AI 套件。等待開發控制檯準備就緒後,單擊開發控制檯。
3. 在開發控制檯左側,選擇 Notebook,在 Notebook 頁面右上角,單擊創建 Notebook 創建新的 Notebook 環境。Notebook 資源需要 CPU:12C,內存:40G,GPU 顯存:24GB。(節點對應規格爲 ecs.gn7i-c16g1.4xlarge[10])
2.3 準備 TensorRT-LLM 環境
1. 構建 Notebook 所需鏡像。
FROM docker.io/nvidia/cuda:12.2.2-cudnn8-runtime-ubuntu22.04
ENV DEBIAN_FRONTEND=noninteractive
RUN apt-get update && apt-get upgrade -y && \
apt-get install -y --no-install-recommends \
libgl1 libglib2.0-0 wget git curl vim \
python3.10 python3-pip python3-dev build-essential \
openmpi-bin libopenmpi-dev jupyter-notebook jupyter
RUN pip3 install tensorrt_llm -U --extra-index-url https://pypi.nvidia.com
RUN pip3 install --upgrade jinja2==3.0.3 pynvml>=11.5.0
RUN rm -rf /var/cache/apt/ && apt-get clean && rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* && \
rm -rf /root/.cache/pip/ && rm -rf /*.whl
WORKDIR /root
RUN git clone https://github.com/NVIDIA/TensorRT-LLM.git --branch v0.7.1
ENTRYPOINT ["sh","-c","jupyter notebook --allow-root --notebook-dir=/root --port=8888 --ip=0.0.0.0 --ServerApp.token=''"]
2. 下載模型,本文以 Baichuan2-7B-Base 爲例。
a.確認 tensorrt_llm 安裝成功
! python3 -c "import tensorrt_llm; print(tensorrt_llm.__version__)"
# 0.7.1
b.安裝 baichuan 依賴
! cd /root/TensorRT-LLM/examples/baichuan
!pip3 install -r requirements.txt
c.下載 Baichuan2-7B-Chat 模型
!yum install git-lfs
!GIT_LFS_SKIP_SMUDGE=1 git clone https://www.modelscope.cn/baichuan-inc/Baichuan2-7B-Chat.git
!cd Baichuan2-7B-Chat/
!git lfs pull
d.將模型編譯爲 TensorRT Engines,權重指定爲 INT8。模型轉換約 5 分鐘。
! cd /root/TensorRT-LLM/examples/baichuan
# Build the Baichuan V2 7B model using a single GPU and apply INT8 weight-only quantization.
! python3 build.py --model_version v2_7b \
--model_dir ./Baichuan2-7B-Chat \
--dtype float16 \
--use_gemm_plugin float16 \
--use_gpt_attention_plugin float16 \
--use_weight_only \
--output_dir ./tmp/baichuan_v2_7b/trt_engines/int8_weight_only/1-gpu/
e.使用構建好的 tensort engines 進行推理
# With INT8 weight-only quantization inference
! python3 ../run.py --input_text "世界上第二高的山峯是哪座?" \
--max_output_len=50 \
--tokenizer_dir=./Baichuan2-7B-Chat \
--engine_dir=./tmp/baichuan_v2_7b/trt_engines/int8_weight_only/1-gpu/
預期輸出:
Input [Text 0]: "世界上第二高的山峯是哪座?"
Output [Text 0 Beam 0]: "世界上第二高的山峯是喀喇崑崙山脈的喬戈裏峯(K2),海拔高度爲8611米。"
2.4 性能測試
1. 使用 TensorRT-LLM 自帶的 benchmark。
向 _allowed_configs dict 中添加 baichuan2_7b_chat 配置,代碼可參考鏈接[11]。
注:0.7.1 版本 benchmark 還未支持 baichuan2 模型,因此需要手動修改下 allowed_configs 配置。
! cd /root/TensorRT-LLM/benchmarks/python
! vim allowed_configs.py
# "baichuan2_7b_chat":
ModelConfig(name="baichuan2_7b_chat",
family="baichuan_7b",
benchmark_type="gpt",
build_config=BuildConfig(
num_layers=32,
num_heads=32,
hidden_size=4096,
vocab_size=125696,
hidden_act='silu',
n_positions=4096,
inter_size=11008,
max_batch_size=128,
max_input_len=512,
max_output_len=200,
builder_opt=None,
)),
運行 benchmark:
! python3 benchmark.py \
-m baichuan2_7b_chat \
--mode plugin \
--engine_dir /root/TensorRT-LLM/examples/baichuan/tmp/baichuan_v2_7b/trt_engines/int8_weight_only/1-gpu \
--batch_size 1 \
--input_output_len "32,50;128,50"
# batch_size 併發度
# input_output_len 輸入輸出的長度,多個測試用例用分號分隔
Expected outputs:
[BENCHMARK] model_name baichuan2_7b_chat world_size 1 num_heads 32 num_kv_heads 32 num_layers 32 hidden_size 4096 vocab_size 125696 precision float16 batch_size 1 input_length 32 output_length 50 gpu_peak_mem(gb) 8.682 build_time(s) 0 tokens_per_sec 60.95 percentile95(ms) 821.977 percentile99(ms) 822.093 latency(ms) 820.348 compute_cap sm86 generation_time(ms) 798.45 total_generated_tokens 49.0 generation_tokens_per_second 61.369
[BENCHMARK] model_name baichuan2_7b_chat world_size 1 num_heads 32 num_kv_heads 32 num_layers 32 hidden_size 4096 vocab_size 125696 precision float16 batch_size 1 input_length 128 output_length 50 gpu_peak_mem(gb) 8.721 build_time(s) 0 tokens_per_sec 59.53 percentile95(ms) 841.708 percentile99(ms) 842.755 latency(ms) 839.852 compute_cap sm86 generation_time(ms) 806.571 total_generated_tokens 49.0 generation_tokens_per_second 60.751
2. 對比 INT8 量化模型與原始模型性能。
原始模型執行命令:
def normal_inference():
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.generation.utils import GenerationConfig
tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto", torch_dtype=torch.bfloat16, trust_remote_code=True)
model.generation_config = GenerationConfig.from_pretrained(model_path)
messages = []
messages.append({"role": "user", "content": prompt})
response = model.chat(tokenizer, messages)
print(response)
INT8 量化模型命令:
def tensorrt_llm_inference():
from subprocess import Popen, PIPE
script = f'''python3 /root/TensorRT-LLM/examples/run.py --input_text \"{prompt}\" \
--max_output_len=50 \
--tokenizer_dir=/root/TensorRT-LLM/examples/baichuan/Baichuan2-7B-Chat \
--engine_dir=/root/TensorRT-LLM/examples/baichuan/tmp/baichuan_v2_7b/trt_engines/int8_weight_only/1-gpu/'''
p = Popen(['sh', '-c', script], stdout=PIPE,
stderr=PIPE)
output, err = p.communicate()
if p.returncode != 0:
print(f"tensorrt_llm_inference() error:{err}")
return
print(output)
TensorRT-LLM 加速方案在採用 INT8 模型量化的情況下,相比於默認的 Baichuan2-7B-Chat 模型,顯存峯值降低了 43.8%,時延降低了 61.1%。
3.相關資料
[1] 參考文獻
[2] SmoothQuant技術
[3] AWQ
[4] GPTQ
[5] Orca: A Distributed Serving System for Transformer-Based Generative Models
[6] MQA(Multi-query Attention)
[7] GQA(Group-query Attention)
[8] 安裝雲原生AI套件
[9] 容器服務管理控制檯
[11]https://github.com/NVIDIA/TensorRT-LLM/blob/12e82e30b0e64b0f7ada0dc5993edd3b05385964/benchmarks/python/allowed_configs.py#L940TensorRT-LLM
[12]參考資料
本文爲阿里雲原創內容,未經允許不得轉載。