作者：顧靜

TensorRT-LLM 如何提升?LLM 模型推理效率

大型語言模型（Large language models,LLM）是基于大量數(shù)據(jù)進行預(yù)訓(xùn)練的超大型深度學(xué)習(xí)模型。底層轉(zhuǎn)換器是一組神經(jīng)網(wǎng)絡(luò)，這些神經(jīng)網(wǎng)絡(luò)由具有 self-attention 的編碼器和解碼器組成。編碼器和解碼器從一系列文本中提取含義，并理解其中的單詞和短語之間的關(guān)系。

當(dāng)前 LLM 模型推理的主要瓶頸是 GPU 顯存資源不足。因此，各類加速框架主要集中于降低 GPU 顯存峰值和提高 GPU 使用率兩大目標(biāo)。

TensorRT-LLM [ 1] 是 NVIDIA 推出的大語言模型（LLM）推理優(yōu)化框架。它提供了一組 Python API 用于定義 LLMs，并且使用最新的優(yōu)化技術(shù)將 LLM 模型轉(zhuǎn)換為 TensorRT Engines，推理時直接使用優(yōu)化后的 TensorRT Engines。

TensorRT-LLM 主要利用以下四項優(yōu)化技術(shù)提升 LLM 模型推理效率。

1. 量化

模型量化技術(shù)是通過降低原始模型的精度來減少模型推理時的 GPU 顯存使用。TensorRT 支持多種模型的多種精度，以下列舉了部分主流模型支持的量化精度。

W8A8 SQ 使用了 SmoothQuant 技術(shù) [ 2] ，在不降低模型推理準(zhǔn)確率的前提下，將模型權(quán)重和激活層都降低為 INT8 精度，顯著減少了 GPU 顯存消耗。

W4A16/W8A16 是指模型權(quán)重為 INT4 或者 INT8，激活層為 FP16 精度。

W4A16 AWQ 以及 W4A16 GPTQ 分別實現(xiàn)了 AWQ [ 3] 和 GPTQ [ 4] 兩篇論文中提到的量化方法。模型權(quán)重為 INT4，激活層為 FP16 精度。

2. In-Flight Batching

傳統(tǒng)的 Batching 技術(shù)為 Static Batching 的，需要等 Batching 中所有序列推理完成后才能進行下一次批次。下圖為一個輸出最大 Token 為 8，Batch size 為 4 的推理過程，使用 Static Batching 技術(shù)。S3 序列在 T5 時刻就已經(jīng)完成推理，但是需要等到 S2 序列在 T8 時刻推理完成后才會處理下一個 sequence，存在明顯的資源浪費。

In-Flight Batching 又名 Continuous Batching 或 iteration-level batching，該技術(shù)可以提升推理吞吐率，降低推理時延。Continuous Batching 處理過程如下，當(dāng) S3 序列處理完成后插入一個新序列 S5 進行處理，提升資源利用率。詳情可參考論文?Orca: A Distributed Serving System for Transformer-Based Generative Models [ 5] 。

3. Attention

Attention 機制用于從序列中提取關(guān)鍵/重要信息，在情感識別、翻譯、問答等任務(wù)中起著至關(guān)重要的作用。Attention 機制按照演進順序可以分為 MHA（Multi-head Attention）、MQA（Multi-query Attention） [ 6] 以及?GQA（Group-query Attention） [ 7] 機制。MQA 和 GQA 都是 MHA 的變種。

MHA 是標(biāo)準(zhǔn)的多頭注意力機制，每個 query 存儲一份 KV，因此需要使用較多的顯存。MQA 所有 query 共享一份 KV，推理時容易丟失一些細節(jié)信息。GQA 將 query 進行分組，組內(nèi)共享一份 KV，可以有效避免 MHA 和 MQA 的問題。

TensorRT-LLM 支持 MHA、MQA 及 GQA 方式，可以在 tensorrt_llm.functional.gpt_attention 查看具體實現(xiàn)。

4. Graph Rewriting

TensorRT-LLM 在將 LLM 模型編譯為 TensorRT Engines 時會對神經(jīng)網(wǎng)絡(luò)進行優(yōu)化，提升執(zhí)行效率。

基于阿里云容器服務(wù) ACK 的實戰(zhàn)體驗

云原生 AI 套件

云原生 AI 套件是阿里云容器服務(wù) ACK 提供的云原生 AI 技術(shù)和產(chǎn)品方案，幫助企業(yè)更快、更高效地落地云原生 AI 系統(tǒng)。

本文將介紹如何基于阿里云容器服務(wù) ACK 云原生 AI 套件，利用 TensorRT-LLM 優(yōu)化 LLM 模型推理。

環(huán)境配置

1. 參考文檔安裝云原生 AI 套件 [ 8] 。

2. 登陸容器服務(wù)管理控制臺 [ 9] ，在左側(cè)導(dǎo)航欄選擇集群 > 應(yīng)用 > 云原生 AI 套件。等待開發(fā)控制臺準(zhǔn)備就緒后，單擊開發(fā)控制臺。

3. 在開發(fā)控制臺左側(cè)，選擇?Notebook，在 Notebook 頁面右上角，單擊創(chuàng)建 Notebook 創(chuàng)建新的 Notebook 環(huán)境。Notebook 資源需要 CPU：12C，內(nèi)存：40G，GPU 顯存：24GB。(節(jié)點對應(yīng)規(guī)格為?ecs.gn7i-c16g1.4xlarge [ 10] )

準(zhǔn)備 TensorRT-LLM 環(huán)境

1. 構(gòu)建 Notebook 所需鏡像。

FROM docker.io/nvidia/cuda:12.2.2-cudnn8-runtime-ubuntu22.04ENV DEBIAN_FRONTEND=noninteractiveRUN 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 jupyterRUN pip3 install tensorrt_llm -U --extra-index-url https://pypi.nvidia.com
RUN pip3 install --upgrade jinja2==3.0.3 pynvml>=11.5.0RUN rm -rf /var/cache/apt/ && apt-get clean && rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* && \rm -rf /root/.cache/pip/ && rm -rf /*.whlWORKDIR /root
RUN git clone https://github.com/NVIDIA/TensorRT-LLM.git --branch v0.7.1ENTRYPOINT ["sh","-c","jupyter notebook --allow-root --notebook-dir=/root --port=8888 --ip=0.0.0.0 --ServerApp.token=''"]

2. 下載模型，本文以 Baichuan2-7B-Base 為例。

a.確認(rèn) 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，權(quán)重指定為 INT8。模型轉(zhuǎn)換約 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.使用構(gòu)建好的 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/

預(yù)期輸出：

Input [Text 0]: "世界上第二高的山峰是哪座？"
Output [Text 0 Beam 0]: "世界上第二高的山峰是喀喇昆侖山脈的喬戈里峰（K2），海拔高度為8611米。"

性能測試

1. 使用 TensorRT-LLM 自帶的 benchmark。

向 _allowed_configs dict 中添加 baichuan2_7b_chat 配置，代碼可參考鏈接 [1****1] 。

🔔 注：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 并發(fā)度
# 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 量化模型與原始模型性能。

原始模型執(zhí)行命令：

def normal_inference():from transformers import AutoModelForCausalLM, AutoTokenizerfrom transformers.generation.utils import GenerationConfigtokenizer = 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, PIPEscript = 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}")returnprint(output)

TensorRT-LLM 加速方案在采用 INT8 模型量化的情況下，相比于默認(rèn)的 Baichuan2-7B-Chat 模型，顯存峰值降低了?*43.8% *，時延降低了?61.1%。

參考文獻：

https://nvidia.github.io/TensorRT-LLM/architecture.html

https://www.anyscale.com/blog/continuous-batching-llm-inference

相關(guān)鏈接：

[1]?TensorRT-LLM

https://github.com/NVIDIA/TensorRT-LLM

[2]?SmoothQuant技術(shù)

https://arxiv.org/abs/2211.10438

[3]?AWQ

https://arxiv.org/abs/2306.00978

[4]?GPTQ

https://arxiv.org/abs/2210.17323

[5]?Orca: A Distributed Serving System for Transformer-Based Generative Models

https://help.aliyun.com/zh/eventbridge/user-guide/transform/?spm=a2c4g.11186623.0.0.501b5750w5RP1Q

[6]?MQA（Multi-query Attention）

https://arxiv.org/abs/1911.02150

[7]?GQA（Group-query Attention）

https://arxiv.org/abs/2307.09288

[8]?安裝云原生AI套件

https://help.aliyun.com/zh/ack/cloud-native-ai-suite/user-guide/deploy-the-cloud-native-ai-suite?spm=a2c4g.11186623.0.0.7e223d92U1aVNf

[9]?容器服務(wù)管理控制臺

https://account.aliyun.com/login/login.htm?oauth_callback=https%3A%2F%2Fcs.console.aliyun.com%2F

[10]?ecs.gn7i-c16g1.4xlarge

https://help.aliyun.com/zh/ecs/user-guide/overview-of-instance-families#gn7i

[11]?鏈接

https://github.com/NVIDIA/TensorRT-LLM/blob/12e82e30b0e64b0f7ada0dc5993edd3b05385964/benchmarks/python/allowed_configs.py#L940