使用docker部署tensorrtllm推理大模型baichuan2-7b

简介

大模型的推理框架，我之前用过vllm和mindie。近期有项目要用tensorrtllm，这里将摸索的过程记录下，特别是遇到的问题。
我的环境是Linux+rt3090

准备docker环境

本次使用docker镜像部署，需要从网上拉取：

docker pull nvcr.io/nvidia/tritonserver:24.08-trtllm-python-py3

The Triton Inference Server的docker的tag可以从这里浏览:https://catalog.ngc.nvidia.com/orgs/nvidia/containers/tritonserver

他要用更新或者更久的版本，都可以从这里获取。

拉取的时间较久，拉完镜像启动本地的docker容器并进入：

#  docker run  -dit --ipc=host --runtime=nvidia --gpus '"device=6,7"' -v /home/xxxxx/docker_shared/:/docker_shared --name triton-llm-test  nvcr.io/nvidia/tritonserver:24.08-trtllm-python-py3

# docker exec -it   triton-llm-test bash

上面的命令中，我将一个磁盘目录映射到docker内的/docker_shared，并指定映射了2张显卡。

我预先下载了baichuan2-7b-chat的模型，在docker内它的目录是/docker_shared/Baichuan2-7B-Chat

转换模型

转换模型的方法可参考：https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html

下载https://github.com/NVIDIA/TensorRT-LLM的代码放到docker中，解压后然后到它的目录中,解压后进入其目录：

cd /docker_shared/TensorRT-LLM-main/examples/baichuan/

提前建好转换模型的输出目录/docker_shared/Baichuan2-7B-trllm/convert_out/ ， 然后执行转换脚本：

python3 convert_checkpoint.py --model_version v2_7b --model_dir /docker_shared/Baichuan2-7B-Chat --output_dir /docker_shared/Baichuan2-7B-trllm/convert_out/ --tp_size 2
>> Total time of converting checkpoints: 00:01:35

# ls /docker_shared/Baichuan2-7B-trllm/convert_out/
config.json  rank0.safetensors  rank1.safetensors

这里我指定了tensor并行度 --tp_size为2，所以生成了2个safetensors文件。

创建引擎目录/docker_shared/Baichuan2-7B-trllm/build_out/，然后将模型编译为tensorrt的引擎文件：

trtllm-build --checkpoint_dir /docker_shared/Baichuan2-7B-trllm/convert_out/  --output_dir /docker_shared/Baichuan2-7B-trllm/build_out/    --gpt_attention_plugin float16 --gemm_plugin float16  --max_batch_size 64

# ls  /docker_shared/Baichuan2-7B-trllm/build_out/
config.json  rank0.engine  rank1.engine

到这里，模型转换的工作就做完了。

简单测试

cd /docker_shared/TensorRT-LLM-main/examples

mpirun -n 2 --allow-run-as-root python3 ./run.py --engine_dir /docker_shared/Baichuan2-7B-trllm/build_out/  --max_output_len 500 --tokenizer_dir /docker_shared/Baichuan2-7B-Chat --input_text "讲一个200字左右的故事"

Output [Text 0 Beam 0]: "
从前，有一个小村庄，村子里住着许多勤劳、善良的人们。他们和睦相处，互相帮助，生活得非常幸福。在这个村子里，有一个年轻人叫小明，他聪明、勇敢，深受村民们的喜爱。

有一天，村子里来了一个陌生的老人，他看起来非常疲惫，小明见状，便主动上前帮助他。老人感激地说：“谢谢你，年轻人。我是从远方来的，我的家乡遭受了一场大旱，庄稼枯萎，水源干涸，我带着唯一的希望——一颗神奇的种子，来寻求帮助。”

小明听了非常同情，他决定帮助老人，于是他们一起回到了村子里。小明将神奇的种子种在了村子的中心，村民们纷纷前来观看。不久，神奇的种子长出了一棵大树，树上结满了晶莹剔透的水珠。村民们欣喜若狂，纷纷取水，不久，村子里的水源恢复了，庄稼也重新焕发生机。

老人感激不已，他告诉小明，这颗神奇的种子叫做“希望之树”，它能给人们带来希望和力量。小明和村民们为了纪念这次奇迹，将这棵树作为村子的象征，代代相传。从此，这个村子变得更加美好，人们过上了幸福美满的生活。"

如果之前设置的-tp_size为1，则不需要mpirun命令，直接执行run.py脚本就行。如果tp_size大于1则需要借助mpirun 并指定-n参数，否则会出错：

RuntimeError: [TensorRT-LLM][ERROR] Assertion failed: MPI size 1 != TP size 2 * PP size 1

准备配置文件

下载https://github.com/triton-inference-server/tensorrtllm_backend的v0.12.0版本，放到docker内解压。

使用模板新建一个目录：

cd  /docker_shared/tensorrtllm_backend-0.12.0
cp all_models/inflight_batcher_llm/ triton_model_repo/ -rf

使用脚本fill_template.py填充配置文件：

python3 tools/fill_template.py -i triton_model_repo/preprocessing/config.pbtxt tokenizer_dir:/docker_shared/Baichuan2-7B-Chat,triton_max_batch_size:64,preprocessing_instance_count:1
python3 tools/fill_template.py -i triton_model_repo/postprocessing/config.pbtxt tokenizer_dir:/docker_shared/Baichuan2-7B-Chat,triton_max_batch_size:64,postprocessing_instance_count:1
python3 tools/fill_template.py -i triton_model_repo/tensorrt_llm_bls/config.pbtxt triton_max_batch_size:64,decoupled_mode:False,bls_instance_count:1,accumulate_tokens:False
python3 tools/fill_template.py -i triton_model_repo/ensemble/config.pbtxt triton_max_batch_size:64
python3 tools/fill_template.py -i triton_model_repo/tensorrt_llm/config.pbtxt triton_backend:tensorrtllm,triton_max_batch_size:64,decoupled_mode:False,max_beam_width:1,engine_dir:/docker_shared/Baichuan2-7B-trllm/build_out/,max_tokens_in_paged_kv_cache:2560,max_attention_window_size:2560,kv_cache_free_gpu_mem_fraction:0.1,exclude_input_in_output:True,enable_kv_cache_reuse:False,batching_strategy:inflight_fused_batching,max_queue_delay_microseconds:0

启动服务并检测测试

cd  /docker_shared/tensorrtllm_backend-0.12.0
python3 scripts/launch_triton_server.py --world_size 2 --model_repo ./triton_model_repo/

启动成功，命令行最后会输出：

I0903 11:20:25.664256 203 grpc_server.cc:2463] "Started GRPCInferenceService at 0.0.0.0:8001"
I0903 11:20:25.664504 203 http_server.cc:4694] "Started HTTPService at 0.0.0.0:8000"
I0903 11:20:25.717022 203 http_server.cc:362] "Started Metrics Service at 0.0.0.0:8002"

简单测试：

测试服务：
curl -X POST localhost:8000/v2/models/ensemble/generate -d '{"text_input": "啥是机器学习啊?", "max_tokens": 500, "bad_words": "", "stop_words": ""}'
返回：
{"batch_index":0,"context_logits":0.0,"cum_log_probs":0.0,"generation_logits":0.0,"model_name":"ensemble","model_version":"1","output_log_probs":[0.0,0.0,0.0,0.0,0.0],"sequence_end":false,"sequence_id":0,"sequence_start":false,"text_output":"\n机器学习（Machine Learning）是人工智能（Artificial Intelligence）的一个子领域，它关注的是如何让计算机系统通过数据和算法自动地学习和改进。简单来说，机器学习就是让计算机从数据中学习规律，并根据这些规律进行预测和决策。\n\n在机器学习中，我们使用算法（Algorithm）来训练模型（Model），使它能够在新的数据上进行预测。这些算法可以分为监督学习（Supervised Learning）、无监督学习（Unsupervised Learning）和强化学习（Reinforcement Learning）等类型。\n\n监督学习是指在训练过程中，我们为模型提供输入数据和对应的正确输出，模型通过学习这些数据之间的映射关系来进行预测。无监督学习则是在没有正确输出的情况下，让模型自己发现数据中的结构和规律。强化学习则是让模型通过与环境的交互，学习如何采取最优行动以获得最大回报。\n\n机器学习已经广泛应用于各个领域，如自然语言处理、计算机视觉、推荐系统、医疗诊断等。随着技术的不断发展，机器学习将继续为我们的生活带来更多便利和创新。"}

查询服务元数据
curl -X GET localhost:8000/v2
返回：
{"name":"triton","version":"2.49.0","extensions":["classification","sequence","model_repository","model_repository(unload_dependents)","schedule_policy","model_configuration","system_shared_memory","cuda_shared_memory","binary_tensor_data","parameters","statistics","trace","logging"]}

错误集锦

Error parsing text-format inference.ModelConfig: 29:17: Expected integer, got: $

I0903 06:55:43.852157 202 model_lifecycle.cc:472] "loading: tensorrt_llm:1"
[libprotobuf ERROR /tmp/tritonbuild/tritonserver/build/_deps/repo-third-party-build/grpc-repo/src/grpc/third_party/protobuf/src/google/protobuf/text_format.cc:337] Error parsing text-format inference.ModelConfig: 29:17: Expected integer, got: $
.......
error: creating server: Internal - failed to load all models

这个报错表示加载配置文件出错，需要借助tools/fill_template.py脚本将各个配置文件进行配置就能解决。

Cuda Runtime (out of memory)

报错信息

[TensorRT-LLM][ERROR] [defaultAllocator.cpp::allocate::31] Error Code 1: Cuda Runtime (out of memory)
[TensorRT-LLM][WARNING] Requested amount of GPU memory (6992953856 bytes) could not be allocated. There may not be enough free memory for allocation to succeed.
[TensorRT-LLM][ERROR] [safeDeserialize.cpp::load::284] Error Code 2: OutOfMemory (Requested size was 6992953856 bytes.)

我之前设置kv_cache_free_gpu_mem_fraction:0.3，结果显存不够，改为0.1就够了。

如果还不够，可以试试量化模型。

AttributeError: ‘BaichuanTokenizer’ object has no attribute ‘vocab’

报错信息

 "Failed to process the request(s) for model 'postprocessing_0_0', message: 
 AttributeError: 'BaichuanTokenizer' object has no attribute 'vocab'\n\nAt:\n 
 /docker_shared/tensorrtllm_backend-main/./triton_model_repo/postprocessing/1/model.py(243): _postprocessing\n
  /docker_shared/tensorrtllm_backend-main/./triton_model_repo/postprocessing/1/model.py(154): execute\n"

我本来下载的是tensorrtllm_backend的最新main分支代码，错误原因在all_models/inflight_batcher_llm/postprocessing/1/model.py的243行，这里使用了tokenizer的vocab成员，但是这个成员不存在：

网上有类似的问题说是要降transformers的版本，我试了没啥用，将tensorrtllm_backend替换为V0.12.0版本就好了。