Streaming-LLM是一種支援無限輸入長度進行 LLM 推理的技術。它利用注意力池來防止注意力視窗轉移時模型崩潰。原始工作是在 PyTorch 中實現的,我們提供SwiftInfer ,這是一個 TensorRT 實現,使 StreamingLLM 更俱生產級。我們的實作是基於最近發布的TensorRT-LLM專案。
我們使用TensorRT-LLM中的 API 來建立模型並運行推理。由於 TensorRT-LLM 的 API 不穩定且變化很快,我們將我們的實作與版本為v0.6.0
的42af740db51d6f11442fd5509ef745a4c043ce51
提交綁定。隨著 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安裝來安裝TensorRT-LLM V0.6.0 。
透過使用 docker,您只需執行以下命令即可安裝 SwiftInfer:
git clone https://github.com/hpcaitech/SwiftInfer.git
cd SwiftInfer
pip install .
如果不使用 docker,我們提供了一個腳本來自動安裝 TensorRT-LLM。
先決條件
請確保您已安裝以下軟體包:
確保 TensorRT >= 9.1.0 和 CUDA 工具包的版本 >= 12.2。
安裝張量:
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 範例,您需要先複製 meta-llama/Llama-2-7b-chat-hf 模型或其他基於 Llama 的變體(例如 lmsys/vicuna-7b-v1.3)的 Hugging Face 儲存庫。然後,您可以執行以下命令來建立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
接下來,您需要下載LMSYS-FastChat提供的MT-Bench資料。
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
您應該期望看到這一代的結果如下:
我們使用原始 PyTorch 版本對 Streaming-LLM 的實作進行了基準測試。我們的實作的基準命令在運行 Llama 範例部分中給出,而原始 PyTorch 實現的基準命令在 torch_streamingllm 資料夾中給出。使用的硬體如下所示:
結果(20輪對話)是:
我們仍在努力進一步提高效能並適應 TensorRT V0.7.1 API。我們還注意到 TensorRT-LLM 在他們的範例中整合了 StreamingLLM,但它似乎更適合單一文字生成而不是多輪對話。
這項工作受到 Streaming-LLM 的啟發,使其可用於生產。在整個開發過程中,我們參考了以下資料,我們希望感謝他們對開源社群和學術界的努力和貢獻。
如果您發現 StreamingLLM 和我們的 TensorRT 實作很有用,請引用我們的儲存庫和肖等人提出的原始工作。來自麻省理工學院漢實驗室。
# our repository
# NOTE: the listed authors have equal contribution
@misc { streamingllmtrt2023 ,
title = { SwiftInfer } ,
year = { 2023 } ,
publisher = { GitHub } ,
journal = { GitHub repository } ,
howpublished = { url{https://github.com/hpcaitech/SwiftInfer} } ,
}
# Xiao's original paper
@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 repo
# as TensorRT-LLM team does not provide a bibtex
# please let us know if there is any change needed
@misc { trtllm2023 ,
title = { TensorRT-LLM } ,
year = { 2023 } ,
publisher = { GitHub } ,
journal = { GitHub repository } ,
howpublished = { url{https://github.com/NVIDIA/TensorRT-LLM} } ,
}