NVIDIATensorRT-LLM是一个专为优化大语言模型(LLM)推理而设计的库。它提供了多种先进的优化技术，包括自定义AttentionKernel、InflightBatching、PagedKVCaching、量化技术(FP8、INT4AWQ、INT8SmoothQuant等)以及更多功能，确保您的NVIDIAGPU能发挥出卓越的推理性能。
我们深知您对易用性的需求，为了让您更快上手，并迅速实现流行模型的高性能推理，我们开发了LLMAPI，通过简洁的指令，您可轻松体验TensorRT-LLM带来的卓越性能！
LLMAPI是一个high-levelPythonAPI，专为LLM推理工作流而设计。以下是一个展示如何使用TinyLlama的简单示例：
fromtensorrt_llmimportLLM,SamplingParamsprompts=["Hello,mynameis","ThepresidentoftheUnitedStatesis","ThecapitalofFranceis","ThefutureofAIis",]sampling_params=SamplingParams(temperature=0.8,top_p=0.95)llm=LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0")outputs=llm.generate(prompts,sampling_params)#Printtheoutputs.foroutputinoutputs:prompt=output.promptgenerated_text=output.outputs[0].textprint(f"Prompt:{prompt!r},Generatedtext:{generated_text!r}")
希望以上示例能帮助您快速入门NVIDIATensorRT-LLMLLMAPI。
当前TensorRT-LLMLLMAPI可支持的模型
一、详细介绍
1.1模型准备
您可以使用LLM(model=)构造函数来灵活切换这些格式。以下各节将详细介绍具体使用方法。
HuggingFaceHub
使用HuggingFaceHub来导入模型非常直观，只需在LLM构造函数中指定模型仓库名称即可：
llm=LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0")
本地HuggingFace模型
gitlfsinstallgitclone
下载完成后，您可以通过以下方式加载模型：
llm=LLM(model=)
请注意，使用此模型需要遵守特定许可条款：
在开始下载之前，请确保同意这些条款并在HuggingFace上完成身份验证：
本地TensorRT-LLM引擎
LLMAPI支持使用TensorRT-LLMEngine，您可以通过以下两种方式构建Engine：
llm=LLM()#Saveenginetolocaldiskllm.save()
1.2使用技巧和故障排除
以下是针对熟悉TensorRT-LLM其他API的用户，在刚开始使用LLMAPI时可能遇到的常见问题及其解决方案：
在使用LLMAPI进行单节点多GPU推理时，无需添加mpirun前缀。您可以直接运行pythonllm_inference_distributed.py来执行多GPU推理。
在使用Slurm管理的节点上遇到挂起或其他问题时，请在启动脚本中添加前缀mpirun-n1--oversubscribe--allow-run-as-root。
示例命令：
mpirun-n1--oversubscribe--allow-run-as-rootpythonllm_inference_distributed.py
由于LLMAPI依赖mpi4py库，为避免mpi4py中的递归生成进程，请将LLM类放在函数中，并在__main__命名空间下保护程序的主入口点。
注意：此限制仅适用于多GPU推理。
二、常见自定义操作
2.1量化
TensorRT-LLM可以通过在LLM实例中设置适当Flags，自动对HuggingFace模型进行量化。例如，要执行Int4AWQ量化，以下代码会触发模型量化。请参考完整的支持的标志列表和可接受的值。
fromtensorrt_llm.llmapiimportQuantConfig,QuantAlgoquant_config=QuantConfig(quant_algo=QuantAlgo.W4A16_AWQ)llm=LLM(,quant_config=quant_config)
2.2采样
SamplingParams可以自定义采样策略以控制LLM生成的响应，如BeamSearch、Temperature和其他参数。
例如，要启用BeamSize为4的BeamSearch，请按如下方式设置Sampling_Params：
fromtensorrt_llm.llmapiimportLLM,SamplingParams,BuildConfigbuild_config=BuildConfig()build_config.max_beam_width=4llm=LLM(,build_config=build_config)#LettheLLMobjectgeneratetextwiththedefaultsamplingstrategy,or#youcancreateaSamplingParamsobjectaswellwithseveralfieldssetmanuallysampling_params=SamplingParams(beam_width=4)#currentlimitation:beam_widthshouldbeequaltomax_beam_widthforoutputinllm.generate(,sampling_params=sampling_params):print(output)
SamplingParams管理并分发字段到C++classes，包括：
SamplingConfig:
OutputConfig:
更多详情请参考class文档：
2.3Build配置
除了上述参数外，您还可以使用build_config类和从trtllm-buildCLI借用的其他参数来自定义构建配置。这些构建配置选项为目标硬件和用例构建Engine提供了灵活性。请参考以下示例：
11m=LLM(,build_config=Buildconfig(max_num_tokens=4096,maxbatchsize=128,max_beam_width=4))
更多详情请参考buildconfig文档：
2.4自定义Runtime
类似于build_config，您也可以使用runtime_config、peft_cache_config，或其他从ExecutorAPI借用的参数来自定义Runtime配置。这些Runtime配置选项在KVcachemanagement、GPUmemoryallocation等方面提供了额外的灵活性。请参考以下示例：
fromtensorrt_llm.llmapiimportLLM,KvCacheConfigllm=LLM(,kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.8))
2.5自定义Tokenizer
默认情况下，LLMAPI使用transformers的AutoTokenizer。您可以在创建LLM对象时传入自己的分词器来覆盖它。请参考以下示例：
llm=LLM(,tokenizer=)
LLM()工作流将使用您的tokenizer。
也可以使用以下代码直接输入tokenID，由于未使用Tokenizers，该代码生成的是不带文本的tokenID。
llm=LLM()foroutputinllm.generate([32,12]):...
更多详细信息欢迎扫描下方二维码查阅以下详细文档。
关于作者
严春伟
NVIDIA性能架构师，目前主要聚焦于大模型推理架构和优化。
张国铭
NVIDIA性能架构师，目前主要从事大模型推理架构和优化。
AdamZheng
NVIDIA产品经理，负责NVIDIAAI平台软件产品管理，目前主要聚焦于大模型推理架构和优化。