ChatGLM3模型基于TensorRT-LLM的高性能部署指南-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00863/article/details/148391789

ChatGLM3模型基于TensorRT-LLM的高性能部署指南

ChatGLM3 ChatGLM3 - 由清华大学和智谱AI联合发布的新一代对话预训练模型，具备强大的语言理解和生成能力。项目地址: https://gitcode.com/gh_mirrors/ch/ChatGLM3

前言

随着大语言模型(LLM)的快速发展，如何高效部署这些模型成为工程实践中的关键挑战。本文将详细介绍如何利用NVIDIA TensorRT-LLM框架来部署ChatGLM3系列模型，实现高性能推理。

TensorRT-LLM简介

TensorRT-LLM是NVIDIA推出的专为大语言模型优化的推理框架，具有以下核心优势：

高性能推理：通过内核融合、算子优化等技术显著提升推理速度
显存优化：支持KV Cache分页管理等技术降低显存占用
量化支持：提供多种量化方案平衡精度与性能
并行计算：支持多GPU并行推理

环境准备

基础环境搭建

首先需要准备TensorRT-LLM的开发环境，推荐使用Docker容器：

# 安装必要的工具
apt-get update && apt-get -y install git git-lfs

# 获取TensorRT-LLM代码
git clone TensorRT-LLM仓库地址
cd TensorRT-LLM

# 切换到v0.7.0版本
git checkout tags/v0.7.0 -b release/0.7.0
git submodule update --init --recursive
git lfs install
git lfs pull

# 构建并运行Docker容器
make -C docker release_build
make -C docker release_run

ChatGLM3模型部署

模型获取

TensorRT-LLM支持ChatGLM3系列多个模型版本：

chatglm3-6b：标准6B参数版本
chatglm3-6b-base：基础版
chatglm3-6b-32k：支持32k上下文长度版本

可通过以下方式获取模型：

cd ./examples/chatglm
pip install -r requirements.txt
apt-get install git-lfs

# 下载所需模型
git clone chatglm3-6b模型地址      chatglm3_6b
git clone chatglm3-6b-base模型地址 chatglm3_6b_base
git clone chatglm3-6b-32k模型地址  chatglm3_6b_32k

引擎构建

使用build.py脚本构建推理引擎，支持多种配置选项：

# FP16精度单GPU引擎
python3 build.py -m chatglm3_6b --output_dir trt_engines/chatglm3_6b/fp16/1-gpu

# 开启FMHA优化
python3 build.py -m chatglm3_6b --enable_context_fmha --output_dir trt_engines/chatglm3_6b/fp16/1-gpu

# INT8量化引擎
python3 build.py -m chatglm3_6b --use_weight_only --output_dir trt_engines/chatglm3_6b/weight_only/1-gpu

# 双GPU并行
python3 build.py -m chatglm3_6b --world_size 2 --output_dir trt_engines/chatglm3_6b/fp16/2-gpu

关键优化选项

插件配置：
- --use_gpt_attention_plugin：配置注意力插件
- --use_gemm_plugin：配置矩阵乘法插件
- --use_rmsnorm_plugin：配置归一化插件
FMHA优化：
- --enable_context_fmha：开启FP16精度的融合多头注意力
- --enable_context_fmha_fp32_acc：开启FP32精度的融合多头注意力
量化选项：
- --use_weight_only：启用权重量化
- --weight_only_precision：选择int4或int8量化
批处理优化：
- --use_inflight_batching：启用动态批处理
- --tokens_per_block：配置KV Cache分块大小

推理实践

基础推理示例

python3 ../run.py --input_text "ChatGLM3-6B相比ChatGLM2-6B有哪些改进?" \
                 --max_output_len 50 \
                 --tokenizer_dir chatglm3_6b \
                 --engine_dir trt_engines/chatglm3_6b/fp16/1-gpu

多GPU推理

mpirun -n 2 \
    python ../run.py --input_text "ChatGLM3-6B相比ChatGLM2-6B有哪些改进?" \
                    --max_output_len 50 \
                    --tokenizer_dir chatglm3_6b \
                    --engine_dir trt_engines/chatglm3_6b/fp16/1-gpu

交互式对话

python3 tensorrt_llm_cli_demo.py --tokenizer_dir chatglm3_6b --engine_dir trt_engines/chatglm3_6b/fp16/1-gpu

文本摘要任务

python3 ../summarize.py --test_trt_llm \
                       --hf_model_dir chatglm3_6b \
                       --engine_dir trt_engines/chatglm3_6b/fp16/1-gpu