在大模型时代，各大公司在陆续推出和优化各自的底座大模型，不断刷新榜单，然而大模型的超大参数给生产部署带来了很大的困难，由此也带来大模型部署框架的蓬勃发展（可以参考之前写的[LLM推理部署（一）：LLM七种推理服务框架总结](http://mp.weixin.qq.com/s?__biz=Mzg3NDIyMzI0Mw==&mid=2247487223&idx=1&sn=571a8677c8a1fa27e137e4f6deccdfd8&chksm=ced54d13f9a2c405d14f39ea77c2fc80f7fa18a377b76074d4662fed79543a5cb6416296a54d&scene=21#wechat_redirect)），本文将介绍英伟达(NVIDIA)在TensorRT基础上针对LLM优化所推出的推理加速引擎TensorRT-LLM（https://nvidia.github.io/TensorRT-LLM/）。
TensorRT-LLM介绍
TensorRT-LLM针对主流LLM使用TensorRT在GPU高效推理的能力，以NVIDIA Triton（https://developer.nvidia.com/nvidia-triton-inference-server）推理服务器集成为后端，可以在python和C++环境下运行，并且支持单GPU和多GPU（ **Tensor并行和Pipeline并行** ）推理，同时也支持量化模式。
TensorRT-LLM为用户提供了一个易于使用的Python API，类似PyTorch API，比如包含einsum、softmax、matmul或view等功能，layer模块（https://github.com/NVIDIA/TensorRT-LLM/blob/release/0.5.0/tensorrt\_llm/layers）提供了一些常用的LLM组件，比如Attention、MLP或者完整的Transformer等。主要的python API如下图所示：
TensorRT-LLM安装
对于windows操作系统可以参考：https://github.com/NVIDIA/TensorRT-LLM/blob/release/0.5.0/windows
TensorRT-LLM必须从源码安装，参考链接如下： https://github.com/NVIDIA/TensorRT-LLM/blob/release/0.5.0/docs/source/installation.md
TensorRT-LLM使用
使用TensorRT-LLM部署大模型大致分为如下三个步骤：
- 下载预训练模型权重；
- 创建大模型的全优化引擎；
- 部署该引擎
下面以BLOOM-560m模型为例介绍使用TensorRT-LLM部署大模型的步骤：
Step0：在docker容器中安装所需要的环境
pip install -r examples/bloom/requirements.txt
git lfs install
Step1：从Huggingface中下载BLOOM-650m模型
cd examples/bloom
rm -rf ./bloom/560M
mkdir -p ./bloom/560M && git clone https://huggingface.co/bigscience/bloom-560m ./bloom/560M
Step2：创建引擎
# Single GPU on BLOOM 560M
python build.py --model_dir ./bloom/560M/ \
--dtype float16 \
--use_gemm_plugin float16 \
--use_gpt_attention_plugin float16 \
--output_dir ./bloom/560M/trt_engines/fp16/1-gpu/
Note：关于参数的细节可以参考https://github.com/NVIDIA/TensorRT-LLM/blob/release/0.5.0/examples/bloom
Step3：运行引擎
python summarize.py --test_trt_llm \
--hf_model_location ./bloom/560M/ \
--data_type fp16 \
--engine_dir ./bloom/560M/trt_engines/fp16/1-gpu/
除了上述example之外，我们来看一下具体API的使用说明：
在模型级别，TensorRT-LLM采用复杂的策略，如内核融合，将其中多个操作合并到单个内核中，以减少启动多个内核的开销。它还利用量化，大大加快了计算速度，减少了内存需求，而不影响模型精度。
import tensorrtllm as trtllm
# Initialize the model
model = trtllm.LargeLanguageModel('./path_to_your_model')
# Apply kernel fusion and quantization
optimization_flags = trtllm.OptimizationFlag.FUSE_OPERATIONS | trtllm.OptimizationFlag.QUANTIZE
optimized_model = model.optimize(flags=optimization_flags)
在运行时级别，TensorRT-LLM具有连续批处理等功能，允许同时计算多个推理请求，有效地提高GPU利用率和吞吐量。分页注意力是另一个新特性，优化了注意力计算过程中的内存使用，这是大型语言模型的一个常见瓶颈。
# Enable in-flight batching and paged attention
runtime_parameters = {
'in_flight_batching': True,
'paged_attention': True
}
# Build the engine with these runtime optimizations
engine = optimized_model.build_engine(runtime_parameters=runtime_parameters)
在当今的数字时代，速度是至关重要的。TensorRT-LLM可与传统方法相比，提供高达8倍的吞吐量。
这种性能上的飞跃在很大程度上归功于in\_flight\_batching。与传统的批处理不同，在传统的批处理中，推理请求是分组处理的(导致单个请求的延迟)，而在线批处理重叠了不同请求的计算，在不影响批大小的情况下大大减少了推理时间。
input_data = [...] # your input data here
results = engine.execute_with_inflight_batching(input_data)
丰富多样的大型语言模型(llm)，每个模型都是为特定任务量身定制的。推理工具的效用因其与各种模型无缝集成的能力而大大增强。TensorRT-LLM在这一领域表现出色，并且提供了广泛的兼容性，从Meta的Llama 1和2到ChatGLM、Falcon、MPT、Baichuan、Starcoder等一系列llm。
import tensorrtllm as trtllm
# Define and load different LLMs
llama_model = trtllm.LargeLanguageModel('./path_to_llama_model')
chatglm_model = trtllm.LargeLanguageModel('./path_to_chatglm_model')
# Build optimized engines for different LLMs
llama_engine = llama_model.build_engine()
chatglm_engine = chatglm_model.build_engine()
部署人工智能的经济方面通常是人工智能驱动项目可行性的决定性因素。除了原始计算性能之外，TensorRT-LLM的设计还具有成本效益，解决了包括直接和间接费用在内的总拥有成本(TCO)问题。通过提高计算效率，TensorRT-LLM减少了对大量硬件资源的依赖，从而降低了能耗。
import tensorrtllm as trtllm
# Initialize the model
model = trtllm.LargeLanguageModel('./path_to_your_model')
# Optimize the model with energy-efficient settings
optimized_model = model.optimize(energy_efficient=True)
# Monitor energy consumption
energy_usage = optimized_model.monitor_energy_usage()
进入大型语言模型(llm)的世界不需要计算机科学博士学位或多年的编程经验。TensorRT-LLM的设计以用户友好为核心。通过其直观的Python API, TensorRT-LLM使LLM优化和推理平民化，使这些先进技术能够为更广泛的受众所使用。
import tensorrtllm as trtllm
# Initialize and load the model
model = trtllm.LargeLanguageModel('./path_to_your_model')
# Perform common operations through easy-to-understand methods
model.optimize()
model.build_engine()
model.execute(input_data)
模型的规模呈指数级增长，管理计算资源至关重要。TensorRT-LLM的量化支持允许使用较低的精度(如FP8)进行计算，TensorRT-LLM在资源消耗、执行速度和模型精度之间实现了良好的平衡。这不仅加快了推理速度，还减少了内存使用，这对于在受限环境中部署大型模型至关重要。
import tensorrtllm as trtllm
# Initialize the model
model = trtllm.LargeLanguageModel('./path_to_your_model')
# Enable quantization
quantized_model = model.enable_quantization(precision='FP8')
# Build and execute the quantized model
engine = quantized_model.build_engine()
result = engine.execute(input_data)
作为NVIDIA官方产品，TensorRT-LLM在构建时考虑了适应性，准备与新兴的LLM生态系统集成。随着新模型架构的出现和现有模型的完善，TensorRT-LLM支持与前沿开发的无缝集成。
import tensorrtllm as trtllm
# Initialize the model
model = trtllm.LargeLanguageModel('./path_to_your_model')
# Update the model with new kernels or architectures
updated_model = model.update_components(new_kernels='./path_to_new_kernels',
new_architectures='./path_to_new_architectures')
# Re-optimize and deploy the updated model
updated_engine = updated_model.build_engine()
TensorRT-LLM支持的GPU类型
目前支持H100、L40S、A100、A30、V100，除此之外的GPU型号可以使用Volta, Turing, Ampere, Hopper和Ada Lovelace架构运行，不过可能会有限制，比如对数据精度的支持就有所不同，如下表所示：
| FP32 | FP16 | BF16 | FP8 | INT8 | INT4 | |
|---|---|---|---|---|---|---|
| Volta (SM70) | Y | Y | N | N | Y | Y |
| Turing (SM75) | Y | Y | N | N | Y | Y |
| Ampere (SM80, SM86) | Y | Y | Y | N | Y | Y |
| Ada-Lovelace (SM89) | Y | Y | Y | Y | Y | Y |
| Hopper (SM90) | Y | Y | Y | Y | Y | Y |
更多可以参考：https://github.com/NVIDIA/TensorRT-LLM/blob/release/0.5.0/docs/source/precision.md
TensorRT-LLM支持的核心技术点
每个大模型对上述技术的支持不同，具体可以参考：https://github.com/NVIDIA/TensorRT-LLM/blob/release/0.5.0/examples
TensorRT-LLM支持的大模型
TensorRT-LLM大模型的推理性能
GPT-J, LLAMA-7B, LLAMA-70B, Falcon-180B模型在H100, L40S 和 A100GPU下测试的性能可以参考：https://github.com/NVIDIA/TensorRT-LLM/blob/release/0.5.0/docs/source/performance.md
TensorRT-LLM的一些高级话题
参考文献：
[1] https://github.com/NVIDIA/TensorRT-LLM