TensorRT-LLM部署大模型 接口调用

### 使用 TensorRT-LLM 部署大型模型 为了在 TensorRT-LLM 中成功部署大型语言模型 (LLMs)，需遵循一系列特定的操作流程来确保最佳性能和兼容性。由于每种模型架构都有所不同，TensorRT 进行的是深层次图级优化，因此并非所有 LLMs 均能开箱即用地得到支持[^1]。 #### 准备工作 安装必要的依赖项之前，确认已具备适当版本的 NVIDIA CUDA 工具包及 cuDNN 库。接着通过 pip 或者 conda 安装 TensorRT 及其 Python 绑定： ```bash pip install nvidia-tensorrt tensorrt_llm ``` #### 加载预训练模型 对于想要加载到 TensorRT-LLM 的预训练模型而言，通常需要先将其转换为目标框架格式。这里以 Hugging Face Transformers 提供的一个 PyTorch 模型为例展示这一过程： ```python from transformers import AutoModelForCausalLM, AutoTokenizer model_name = "mistralai/Mistral-7B-Instruct-v0.1" tokenizer = AutoTokenizer.from_pretrained(model_name) pytorch_model = AutoModelForCausalLM.from_pretrained(model_name) ``` #### 构建 TensorRT 引擎 一旦拥有合适的输入模型文件路径，就可以利用 `tensorrt_llm` 创建对应的 TensorRT 引擎实例了。下面是一段简单的代码片段用于说明此操作： ```python import tensorrt as trt from tensorrt_llm.models import load_from_hf_hub engine_path = 'path/to/save/engine.trt' with open(engine_path, mode='wb') as f: with trt.Builder(trt.Logger()) as builder, \ builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) as network, \ trt.OnnxParser(network, trt.Logger()) as parser: model_config = load_from_hf_hub(model_name=model_name) # 将PyTorch模型导出为ONNX格式 onnx_file_path = './temp_model.onnx' torch.onnx.export( pytorch_model, args=(torch.randint(high=50256, size=(1, 8), dtype=torch.int64).cuda(),), f=onnx_file_path, input_names=['input_ids'], output_names=["logits"], opset_version=13, ) # 解析 ONNX 文件并配置网络参数 with open(onnx_file_path, 'rb') as model_file: if not parser.parse(model_file.read()): raise RuntimeError('Failed to parse the ONNX file') config = builder.create_builder_config() profile = builder.create_optimization_profile() min_input_shape = {"input_ids": [1, 1]} opt_input_shape = {"input_ids": [1, 8]} max_input_shape = {"input_ids": [1, 64]} profile.set_shape("input_ids", min=min_input_shape["input_ids"],\ opt=opt_input_shape["input_ids"],max=max_input_shape["input_ids"]) config.add_optimization_profile(profile) serialized_engine = builder.build_serialized_network(network, config) f.write(serialized_engine) ``` 这段脚本展示了如何创建一个基于给定 ONNX 模型的新 TensorRT 引擎，并保存至磁盘以便后续使用。注意，在实际应用中可能还需要调整更多选项来适应具体需求。 #### 执行推理请求 最后一步就是编写客户端应用程序向已经准备好的 TensorRT 引擎发送查询请求。这可以通过读取先前生成的序列化引擎数据实现，并建立相应的上下文环境来进行预测计算： ```python def infer(input_text): inputs = tokenizer.encode_plus(prompt=input_text, return_tensors="pt").to('cuda') runtime = trt.Runtime(trt.Logger()) with open(engine_path, 'rb') as f: engine_data = f.read() engine = runtime.deserialize_cuda_engine(engine_data) context = engine.create_execution_context() bindings = [ None, # Input binding index is set to None because we will provide it directly. cuda.mem_alloc(inputs['input_ids'].nelement() * inputs['input_ids'].element_size()), ] stream = cuda.Stream() outputs = np.empty([1], dtype=np.float32) # Transfer data from CPU to GPU and execute inference asynchronously. cuda.memcpy_htod_async(bindings[1], inputs['input_ids'], stream) context.execute_v2(bindings) cuda.memcpy_dtoh_async(outputs, bindings[1], stream) stream.synchronize() result_tensor = torch.tensor(outputs).reshape_as(inputs['input_ids']) generated_tokens = tokenizer.decode(result_tensor.tolist()[0]) return generated_tokens ``` 上述函数接收一段文本作为输入，经过编码处理后传递给 TensorRT 引擎完成推理任务，最终返回解码后的输出字符串表示形式的结果。