TensorRT-LLM自动部署环境搭建与开发指南

TensorRT-LLM自动部署环境搭建与开发指南

TensorRT-LLM TensorRT-LLM provides users with an easy-to-use Python API to define Large Language Models (LLMs) and build TensorRT engines that contain state-of-the-art optimizations to perform inference efficiently on NVIDIA GPUs. TensorRT-LLM also contains components to create Python and C++ runtimes that execute those TensorRT engines. 项目地址: https://gitcode.com/gh_mirrors/te/TensorRT-LLM

前言

TensorRT-LLM是NVIDIA推出的高性能推理引擎，专为大型语言模型(LLM)优化设计。其中的自动部署模块(AutoDeploy)提供了便捷的模型部署方案。本文将详细介绍如何搭建开发环境并进行代码贡献。

环境准备

1. 创建虚拟环境

推荐使用conda管理Python环境，特别是处理非pip依赖项(如openmpi)时更为方便。

conda create -y -n auto pip python=3.12
conda activate auto

2. 安装MPI依赖

分布式推理需要MPI支持，通过conda安装相关组件：

conda install -y -c conda-forge mpi4py openmpi

设置必要的环境变量：

export LD_LIBRARY_PATH=$CONDA_PREFIX/lib:$LD_LIBRARY_PATH
export OMPI_MCA_opal_cuda_support=true

建议将这些配置添加到shell启动文件中(~/.bashrc或~/.zshrc)。

3. 安装TensorRT-LLM

使用预编译的wheel包可以简化安装过程：

export TRTLLM_PRECOMPILED_LOCATION=https://pypi.nvidia.com/tensorrt-llm/tensorrt_llm-0.18.0.dev2025021800-cp312-cp312-linux_x86_64.whl
pip install -e ".[devel]"

开发流程

代码规范检查

项目使用pre-commit进行代码质量管控：

1. 安装pre-commit工具：

pip install pre-commit
pre-commit install

2. 提交代码时自动运行检查，如需跳过可使用：

git commit -n ...

3. 手动运行检查：

pre-commit run --all-files

测试验证

项目采用pytest测试框架：

运行完整测试套件：

pytest tests/_torch/autodeploy

调试技巧

调试时建议设置world_size=0避免子进程干扰：

1. 使用标准Python调试器
2. 在VSCode中可利用项目提供的推荐配置

验证安装

运行示例脚本验证环境配置：

cd examples/auto_deploy
python build_and_run_ad.py --config '{"model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0"}'

开发建议

1. 定期更新预编译wheel包以保持兼容性
2. 开发前确保MPI环境变量正确设置
3. 充分利用VSCode提供的开发配置
4. 调试时简化环境配置(如单进程模式)

通过以上步骤，开发者可以快速搭建TensorRT-LLM自动部署模块的开发环境，并按照规范流程进行代码开发和测试。

TensorRT-LLM TensorRT-LLM provides users with an easy-to-use Python API to define Large Language Models (LLMs) and build TensorRT engines that contain state-of-the-art optimizations to perform inference efficiently on NVIDIA GPUs. TensorRT-LLM also contains components to create Python and C++ runtimes that execute those TensorRT engines. 项目地址: https://gitcode.com/gh_mirrors/te/TensorRT-LLM