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. It also includes a backend for integration with the NVIDIA Triton Inference Server. Models built with TensorRT-LLM can be executed on a wide range of configurations going from a single GPU to multiple nodes with multiple GPUs (using Tensor Parallelism).
The Python API of TensorRT-LLM is architectured to look similar to the
PyTorch API. It provides users with a
functional module containing functions like
einsum
, softmax
, matmul
or view
. The layer
module bundles useful building blocks to assemble LLMs; like an Attention
block, a MLP
or the entire Transformer
layer. Model-specific components,
like GPTAttention
or BertAttention
, can be found in the
model module.
TensorRT-LLM provides users with predefined models that can easily be modified
and extended. The current version of TensorRT-LLM supports
BERT,
GPT,
NVIDIA GPT-2B,
GPT-J,
LLaMA,
OPT,
SantaCoder
and
StarCoder.
To maximize performance and reduce memory footprint, TensorRT-LLM allows the
models to be executed using different quantization modes (see
examples/gpt
for concrete examples). TensorRT-LLM supports
INT4 or INT8 weights (and FP16 activations; a.k.a. INT4/INT8 weight-only) as
well as a complete implementation of the
SmoothQuant technique.
For a more detailed presentation of the software architecture and the key concepts used in TensorRT-LLM, we recommend you to read the following document.
TensorRT-LLM contains Python and C++ components, and must be compiled from source to be used. TensorRT-LLM is dependent on the latest versions of TensorRT and Polygraphy which are distributed separately, and should be copied into this repository.
We recommend that you use a Docker container to build and run TensorRT-LLM. Instructions to install an environment to run Docker containers for the NVIDIA platform can be found here.
Make sure you have fetched all the dependencies before compiling TensorRT-LLM:
git submodule update --init --recursive
Use the following command to create a Docker image for development:
make -C docker build
This will create a docker image for development of TensorRT-LLM and tag it locally with tensorrt_llm/devel:latest
.
To run the container, use the following command:
make -C docker run
If you prefer to work with your own user account in that container instead of root
, include the option LOCAL_USER=1
in the command above like so:
make -C docker run LOCAL_USER=1
make
On systems without GNU make
or shell support, you can build the Docker image for development as follows:
docker build --pull \
--target devel \
--file docker/Dockerfile.multi \
--tag tensorrt_llm/devel:latest \
.
Then run the container by issuing the following command:
docker run --rm -it \
--ipc=host --ulimit memlock=-1 --ulimit stack=67108864 --gpus=all \
--volume ${PWD}:/code/tensorrt_llm \
--workdir /code/tensorrt_llm \
tensorrt_llm/devel:latest
# To build the TensorRT-LLM code.
python3 ./scripts/build_wheel.py --trt_root /usr/local/tensorrt
# Deploy TensorRT-LLM in your environment.
pip install ./build/tensorrt_llm*.whl
By default, build_wheel.py
enables incremental builds. To clean the build
directory, add the --clean
option:
python3 ./scripts/build_wheel.py --clean --trt_root /usr/local/tensorrt
The steps of creating a Docker image for development, building the wheel and installing it inside the container can be executed in a single command:
make -C docker release_build
You can optionally append CUDA_ARCHS="<list of architectures in CMake format>"
to specify which architectures should
be supported by the wheel. Once the image is built, run it in a Docker container with:
make -C docker release_run
Append LOCAL_USER=1
to this command for switching to your local user account instead of root
inside the container.
The examples of TensorRT-LLM are installed in directory /app/tensorrt_llm/examples
.
Specific CUDA architectures may be passed as an argument to
build_wheel.py
. The script accepts a single
argument taking a semicolon separated list of CUDA architecture specifications
compatible with CUDA_ARCHITECTURES in CMake. For instance, to build for
compute capabilities 8.0 and 8.6, call build_wheel.py
like so:
python3 ./scripts/build_wheel.py --cuda_architectures "80-real;86-real"
Running build_wheel.py
will also compile the library containing the C++
runtime of TensorRT-LLM. If Python support and torch
modules are not
required, the script provides the option --cpp_only
which restricts the build
to the C++ runtime only:
python3 ./scripts/build_wheel.py --cuda_architectures "80-real;86-real" --cpp_only --clean
This is particularly useful to avoid linking problems which may be introduced
by particular versions of torch
related to the dual ABI support of
GCC. The
option --clean
will remove the build directory before building. The default
build directory is cpp/build
, which may be overridden using the option
--build_dir
. Run build_wheel.py --help
for an overview of all supported
options.
Clients may choose to link against the shared or the static version of the library. These libraries can be found in the following locations:
cpp/build/tensorrt_llm/libtensorrt_llm.so
cpp/build/tensorrt_llm/libtensorrt_llm_static.a
In addition, one needs to link against the library containing the LLM plugins for TensorRT available here:
cpp/build/tensorrt_llm/plugins/libnvinfer_plugin_tensorrt_llm.so
Add the following directories to your project include paths
cpp
cpp/include
Only header files contained in cpp/include
are part of the supported API and
may be directly included. Other headers contained under cpp
should not be
included directly since they might change in future versions.
For examples of how to use the C++ runtime, see the unit tests in gptSessionTest.cpp and the related CMakeLists.txt file.
It's recommended to add options –shm-size=1g –ulimit memlock=-1
to the
docker or nvidia-docker run command. Otherwise you may see NCCL errors when
running multiple GPU inferences. See
https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/troubleshooting.html#errors
for details.
If you encounter
NVIDIA H100 PCIe with CUDA capability sm_90 is not compatible with the current PyTorch installation. The current PyTorch install supports CUDA capabilities sm_37 sm_50 sm_60 sm_70 sm_75 sm_80 sm_86.
when building engines, you need to install the preview version of PyTorch that corresponds to your CUDA version. As an example, for CUDA 12.1, use:
pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu121
August 2023
July 2023
[batch_size, num_heads, seq_length, dim_per_head]
,functional.py
and documentation in docs
)June 2023
May 2023
此处可能存在不合适展示的内容,页面不予展示。您可通过相关编辑功能自查并修改。
如您确认内容无涉及 不当用语 / 纯广告导流 / 暴力 / 低俗色情 / 侵权 / 盗版 / 虚假 / 无价值内容或违法国家有关法律法规的内容,可点击提交进行申诉,我们将尽快为您处理。