用TensorRT-LLM跑通BLOOM模型_tensorrt-llm bloom7

零、参考资料

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
1
2
3
4
5
6

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

1、配置python环境

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

2、从 HuggingFace 下载模型权重

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

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/
1
2
3
4
5
6
7
8
9

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/
1
2
3
4