🚀 SwiftInfer
🔗 目录
📌 概述
Streaming-LLM 是一种支持无限输入长度的LLM推理技术。它利用 Attention Sink 来防止模型在注意窗口移动时崩溃。原始工作是在 PyTorch 中实现的,我们提供了 SwiftInfer,这是一个 TensorRT 实现,使 StreamingLLM 更具生产级别。我们的实现基于最近发布的 TensorRT-LLM 项目。
🚗 快速开始
🛠 安装
我们使用 TensorRT-LLM 中的 API 来构建模型并进行推理。由于 TensorRT-LLM 的 API 处于不稳定状态并且变化迅速,我们绑定了 42af740db51d6f11442fd5509ef745a4c043ce51
提交的版本 v0.6.0
。随着 TensorRT-LLM 的 API 趋于稳定,我们可能会升级此存储库。
如果你已经构建了 TensorRT-LLM V0.6.0,只需运行:
git clone https://github.com/hpcaitech/SwiftInfer.git
cd SwiftInfer
pip install .
否则,你应先安装 TensorRT-LLM。
使用 Docker 安装 TensorRT-LLM
如果使用 docker,你可以按照 TensorRT-LLM 安装文档 来安装 TensorRT-LLM V0.6.0。
通过使用 docker,你可以只需运行以下命令来安装 SwiftInfer:
git clone https://github.com/hpcaitech/SwiftInfer.git
cd SwiftInfer
pip install .
不使用 Docker 安装 TensorRT-LLM
如果不使用 docker,我们提供了一个脚本来自动安装 TensorRT-LLM。
先决条件
请确保你已经安装了以下软件包:
- python
- 构建工具,包括 gcc/g++,make,cmake
- CUDA toolkit
- cuDNN
- NCCL
- TensorRT
- PyTorch
确保 TensorRT 的版本 >= 9.1.0 并且 CUDA toolkit 的版本 >= 12.2。
安装 tensorrt:
ARCH=$(uname -m)
if [ "$ARCH" = "arm64" ];then ARCH="aarch64";fi
if [ "$ARCH" = "amd64" ];then ARCH="x86_64";fi
if [ "$ARCH" = "aarch64" ];then OS="ubuntu-22.04"; else OS="linux";fi
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/secure/9.1.0/tars/tensorrt-9.1.0.4.$OS.$ARCH-gnu.cuda-12.2.tar.gz
tar xzvf tensorrt-9.1.0.4.linux.x86_64-gnu.cuda-12.2.tar.gz
PY_VERSION=$(python -c 'import sys; print(".".join(map(str, sys.version_info[0:2])))')
PARSED_PY_VERSION=$(echo "${PY_VERSION//./}")
pip install TensorRT-9.1.0.4/python/tensorrt-*-cp${PARSED_PY_VERSION}-*.whl
export TRT_ROOT=$(realpath TensorRT-9.1.0.4)
下载 nccl,请访问 NCCL 下载页面。
下载 cudnn,请访问 cuDNN 下载页面。
命令
在运行以下命令之前,请确保已正确设置 nvcc
。检查方式如下:
nvcc --version
安装 TensorRT-LLM 和 SwiftInfer,运行:
git clone https://github.com/hpcaitech/SwiftInfer.git
cd SwiftInfer
TRT_ROOT=xxx NCCL_ROOT=xxx CUDNN_ROOT=xxx pip install .
🕹 运行 Llama 示例
要运行 Llama 示例,你需要首先克隆 Hugging Face 仓库中的 meta-llama/Llama-2-7b-chat-hf 模型或其他基于 Llama 的变体,例如 lmsys/vicuna-7b-v1.3。然后,你可以运行以下命令来构建 TensorRT 引擎。你需要将 <model-dir>
替换为 Llama 模型的实际路径。
cd examples/llama
python build.py \
--model_dir <model-dir> \
--dtype float16 \
--enable_context_fmha \
--use_gemm_plugin float16 \
--max_input_len 2048 \
--max_output_len 1024 \
--output_dir ./output/7B-streaming-8k-1k-4-2000/trt_engines/fp16/1-gpu/ \
--max_batch_size 1
接下来,你需要下载 MT-Bench 数据,该数据由 LMSYS-FastChat 提供。
mkdir mt_bench_data
wget -P ./mt_bench_data https://raw.githubusercontent.com/lm-sys/FastChat/main/fastchat/llm_judge/data/mt_bench/question.jsonl
最后,你可以使用以下命令运行 Llama 示例。
❗️❗️❗️ 在此之前,请注意:
only_n_first
参数用于控制需要评估的样本数量。如果你想评估所有样本,请删除此参数。
python ../run_conversation.py \
--max_input_length 2048 \
--max_output_len 1024 \
--tokenizer_dir <model-dir> \
--engine_dir ./output/7B-streaming-8k-1k-4-2000/trt_engines/fp16/1-gpu/ \
--input_file ./mt_bench_data/question.jsonl \
--streaming_llm_start_size 4 \
--only_n_first 5
你应该会看到如下的生成输出:
⚖️ 基准测试
我们已经对我们的 Streaming-LLM 实现进行了基准测试,并与 原始 PyTorch 版本 进行了比较。我们的实现的基准测试命令列在 运行 Llama 示例 部分,而原始 PyTorch 实现的基准测试命令列在 torch_streamingllm 文件夹中。使用的硬件如下:
- GPU:Nvidia H800 (80GB)
- CPU:Intel(R) Xeon(R) Platinum 8468
- RAM:2TB
结果(20 轮对话)如下:
我们仍在努力进一步提升性能并适应 TensorRT V0.7.1 的 API。我们还注意到,TensorRT-LLM 已经在其 示例 中集成了 StreamingLLM,但它似乎更适合单一文本生成而不是多轮对话。
🗺 路线图
- 基于 TRT-LLM API 的 Streaming-LLM 注意力实现
- KV 缓存适配
- 提前停止适配
- 连续张量修复
- 基于多轮对话的 Llama 示例
📃 致谢
这个工作受 Streaming-LLM 的启发,使其可以在生产中使用。在开发过程中,我们参考了以下材料,并希望感谢他们对开源社区和学术界的贡献。
📝 引用
如果你发现 StreamingLLM 和我们的 TensorRT 实现有用,请引用我们的存储库和 MIT Han Lab 的 Xiao 等人 提出的原始工作。
# 我们的存储库
# 注意:列出的作者贡献相同
@misc{streamingllmtrt2023,
title = {SwiftInfer},
year = {2023},
publisher = {GitHub},
journal = {GitHub repository},
howpublished = {\url{https://github.com/hpcaitech/SwiftInfer}},
}
# Xiao 的原始论文
@article{xiao2023streamingllm,
title={Efficient Streaming Language Models with Attention Sinks},
author={Xiao, Guangxuan and Tian, Yuandong and Chen, Beidi and Han, Song and Lewis, Mike},
journal={arXiv},
year={2023}
}
# TensorRT-LLM 仓库
# 因为 TensorRT-LLM 团队没有提供 bibtex
# 如果有任何更改需求请告诉我们
@misc{trtllm2023,
title = {TensorRT-LLM},
year = {2023},
publisher = {GitHub},
journal = {GitHub repository},
howpublished = {\url{https://github.com/NVIDIA/TensorRT-LLM}},
}