用TensorRT-LLM跑通BLOOM模型

本文主要是介绍用TensorRT-LLM跑通BLOOM模型，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

零、参考资料

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
make -C docker release_build
make -C docker release_run LOCAL_USER=1

二、在docker镜像中配置并跑Bloom模型

1、配置python环境

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

2、从 HuggingFace 下载模型权重

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

3、将Hugging Face提供的BLOOM模型格式转换成TensorRT格式

# Single GPU on BLOOM 560M
python convert_checkpoint.py --model_dir ./bloom/560M/ \--dtype float16 \--output_dir ./bloom/560M/trt_ckpt/fp16/1-gpu/
# May need to add trtllm-build to PATH, export PATH=/usr/local/bin:$PATH
trtllm-build --checkpoint_dir ./bloom/560M/trt_ckpt/fp16/1-gpu/ \--gemm_plugin float16 \--gpt_attention_plugin float16 \--output_dir ./bloom/560M/trt_engines/fp16/1-gpu/

4、对 CNN Daily 数据集中的文章进行摘要

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

这篇关于用TensorRT-LLM跑通BLOOM模型的文章就介绍到这儿，希望我们推荐的文章对编程师们有所帮助！

---