用TensorRT-LLM加速大模型推理的5个技巧

最近我在公司折腾大模型推理加速，用了NVIDIA新出的TensorRT-LLM框架。这玩意儿真不是盖的，把原来跑得像蜗牛的模型一下子提速了好几倍。不过要玩转它也没那么容易，踩了不少坑才搞明白。

今儿跟大伙儿分享下我的心得体会，希望能给你们省点儿力气。

1. 精度调教很关键

TensorRT-LLM支持FP16和INT8量化，但不是所有层都适合降低精度。我一开始图省事，把整个模型都量化成INT8，结果精度掉得厉害。后来才发现，得根据不同层的特点来决定用啥精度。

比如这样:

import tensorrt_llm as trt_llm

builder = trt_llm.Builder()
network = builder.create_network()

# 对不同层设置不同精度
config = builder.create_builder_config()
config.set_flag(trt_llm.BuilderFlag.FP16)
config.set_flag(trt_llm.BuilderFlag.INT8)

# 敏感层保持FP16精度
config.set_precision_of_layer("attention", trt_llm.DataType.FLOAT16)
config.set_precision_of_layer("feed_forward", trt_llm.DataType.FLOAT16)

# 其他层可以用INT8
config.set_precision_of_layer("embedding", trt_llm.DataType.INT8)

这样搞既能保证关键层的精度，又能尽量压缩模型大小，简直就是鱼和熊掌兼得。

2. 合理设置batch size

batch size设多大是个技术活。太小了GPU利用率低，太大了显存又吃不消。我发现用动态batch size是个好主意，能根据实际负载自动调整。

# 设置动态batch size
profile = builder.create_optimization_profile()
profile.set_shape("input_ids", (1, 1), (16, 512), (64, 1024))
config.add_optimization_profile(profile)

温馨提示:动态batch size虽好，但inference时间可能会有波动。对延迟敏感的场景要谨慎使用。

3. 利用Tensor并行

要是你手头有多张GPU，那可得好好利用起来。TensorRT-LLM支持Tensor并行，能把大模型切片到不同GPU上。

import tensorrt_llm as trt_llm
from tensorrt_llm.builder import Builder
from tensorrt_llm.layers import RowParallelLinear, ColumnParallelLinear

builder = Builder()
network = builder.create_network()

# 设置并行度
world_size = 4
rank = 0# 当前GPU的rank

# 使用并行层
fc1 = ColumnParallelLinear(network, 1024, 4096, bias=True, gather_output=False, parallel_mode="tensor")
fc2 = RowParallelLinear(network, 4096, 1024, bias=True, input_is_parallel=True, parallel_mode="tensor")

不过要注意,并行度设太高反而可能降低效率。得根据模型大小和GPU数量来权衡。

4. 合理安排内存

TensorRT-LLM默认会把整个模型加载到GPU显存里，大模型很容易就把显存撑爆。

这时候就得动点脑筋，比如用显存换出技术：

import tensorrt_llm as trt_llm

config = trt_llm.BuilderConfig()

# 启用显存换出
config.set_memory_pool_limit(trt_llm.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB workspace
config.set_flag(trt_llm.BuilderFlag.ENABLE_PAGED_MEMORY)

这样一来,TensorRT-LLM就会根据需要动态地把不用的数据换出到CPU内存，省下宝贵的显存。

当然了,这么搞多少会影响些性能,得自己权衡。

5. 善用profiling工具

要真正榨干GPU的性能,光凭感觉可不行。得用专业工具来分析瓶颈。NVIDIA提供了不少好东西,比如NSight Systems就挺好用:

nsys profile python your_inference_script.py

这玩意儿能给你生成一份详细的timeline,清清楚楚地显示每个CUDA kernel的执行时间。找出那些最耗时的操作,针对性地优化,效果立竿见影。

不过话说回来,这些技巧也不是万能的。每个模型情况不一样,得根据实际情况来。多尝试,多总结,慢慢就能找到最适合自己模型的优化方案了。反正我是觉得,跟TensorRT-LLM较劲的过程挺有意思,每次看到推理速度蹭蹭往上涨,那感觉简直不要太爽。

醋醋百科网

Good Luck To You!