本文主要是介绍用TensorRT-LLM跑通chatGLM2_6B模型,希望对大家解决编程问题提供一定的参考价值,需要的开发者们随着小编来一起学习吧!
零、参考资料
NVIDIA官网
THUDM的Github
NVIDIA的Github
一、构建 TensorRT-LLM的docker镜像
git lfs install
git clone https://github.com/NVIDIA/TensorRT-LLM.git
cd TensorRT-LLM
git submodule update --init --recursive
sudo make -C docker release_build
sudo make -C docker release_run
二、在docker镜像中配置并跑GLM模型
1、配置python环境
cd ./examples/chatglm
pip install -r requirements.txt
apt-get update
apt-get install git-lfs
2、从 HuggingFace 下载模型权重
git clone https://huggingface.co/THUDM/chatglm2-6b chatglm2_6b
需要等一段时间
3、将Hugging Face提供的GLM模型转换成TensorRT格式
python3 convert_checkpoint.py --model_dir chatglm2_6b --output_dir trt_ckpt/chatglm2_6b/fp16/1-gpu
4、构建 TensorRT 引擎
# ChatGLM2-6B: single-gpu engine with dtype float16, GPT Attention plugin, Gemm plugin
trtllm-build --checkpoint_dir trt_ckpt/chatglm2_6b/fp16/1-gpu --gemm_plugin float16 --output_dir trt_engines/chatglm2_6b/fp16/1-gpu
5、运行例子
# Run the default engine of ChatGLM2-6B on single GPU, other model name is available if built.
python3 ../run.py --input_text "请输入你的问题" \--max_output_len 50 \--tokenizer_dir chatglm2_6b \--engine_dir trt_engines/chatglm2_6b/fp16/1-gpu
这篇关于用TensorRT-LLM跑通chatGLM2_6B模型的文章就介绍到这儿,希望我们推荐的文章对编程师们有所帮助!