背景介紹

大語言模型正以其驚人的新能力推動人工智能的發(fā)展，擴(kuò)大其應(yīng)用范圍。然而，由于這類模型具有龐大的參數(shù)規(guī)模，部署和推理的難度和成本極高，這一挑戰(zhàn)一直困擾著 AI 領(lǐng)域。此外，當(dāng)前存在大量支持模型部署和推理的框架和工具，如 ModelScope 的 Model Pipelines API 和 HuggingFace 的 Text Generation Inference 等，各自都有其獨特的特點和優(yōu)勢。然而，這些工具往往未能充分發(fā)揮 GPU 的性能。

為了解決這些問題，NVIDIA 推出了一種全新的解決方案——TensorRT-LLM。這是一款高度優(yōu)化的開源計算框架，它將NVIDIA TensorRT的深度學(xué)習(xí)編譯器、FasterTransformer 的優(yōu)化內(nèi)核、預(yù)處理和后處理，以及多 GPU/多節(jié)點通信等功能封裝在一個簡單的開源 Python/C++ API 中，同時與硬件進(jìn)行了一體化優(yōu)化，形成了一種產(chǎn)品級的大模型推理解決方案。NVIDIA TensorRT-LLM具有多項突出的特性，包括支持新的 FP8 數(shù)據(jù)格式，這使得模型可以在更低的精度下運(yùn)行，從而減少內(nèi)存消耗，同時保持模型的準(zhǔn)確性。它還支持一種名為“In-flight batching”的新調(diào)度技術(shù)，可以更有效地處理動態(tài)負(fù)載，提高 GPU 利用率。

此外，TensorRT-LLM 還支持模型的并行化和分布式推理，利用張量并行性進(jìn)行模型并行化，使模型可以在多個 GPU 之間并行運(yùn)行，從而實現(xiàn)大型模型的高效推理。最重要的是，TensorRT-LLM 極大地簡化了開發(fā)流程，使得開發(fā)者無需深入了解底層的技術(shù)細(xì)節(jié)，也無需編寫復(fù)雜的 CUDA/C++ 代碼。它提供了一個易用、開源和模塊化的應(yīng)用編程接口，使開發(fā)者能夠輕松定義、優(yōu)化和執(zhí)行新的大語言模型架構(gòu)和增強(qiáng)功能?？偟膩碚f，TensorRT-LLM 讓用戶可以專注于模型的設(shè)計和優(yōu)化，而將底層的性能優(yōu)化工作交給 TensorRT 來完成，大大提高了開發(fā)效率和生產(chǎn)效率，真正實現(xiàn)了大模型推理的易用性和高效性。

阿里云的通義千問開源模型 Qwen-7B，擁有 70 億參數(shù)，在一系列全方位的評估中展示了其在自然語言理解與生成、數(shù)學(xué)問題求解、代碼生成等領(lǐng)域的優(yōu)秀能力。這些評估涵蓋了多個數(shù)據(jù)集，包括 MMLU、C-Eval、GSM8K、HumanEval 以及 WMT22 等。在這些評測中，Qwen-7B 不僅超越了同等規(guī)模的其他大語言模型，甚至在某些方面超過了參數(shù)規(guī)模更大的模型。因此，對于 TensorRT-LLM 來說，支持 Qwen 系列模型具有重要的意義。

開發(fā)與優(yōu)化過程

我們是社區(qū)開發(fā)者，通過阿里云天池舉辦的NVIDIA TensorRT Hackathon 2023接觸到了 NVIDIA TensorRT-LLM，并為它貢獻(xiàn)了代碼。TensorRT-LLM 已開源（https://github.com/NVIDIA/TensorRT-LLM），包含了我們開發(fā)的 Qwen-7B 模型。以下是我們的開發(fā)記錄，供其他開發(fā)者參考。

基礎(chǔ)功能支持

1. 首先我們初步分析了 examples/llama 代碼，以深化對 trt-llm 基本流程的理解。在 llama 項目的 weight.py 中，存在一個 load_from_meta_llama 函數(shù)，該函數(shù)包含 tensorrt_llm.models.LLaMAForCausalLM，此部分定義了 TensorRT 的模型結(jié)構(gòu)。復(fù)制 examples/llama 并將其重命名為 examples/qwen，同時將 LLaMAForCausalLM 復(fù)制并創(chuàng)建新的 mode.py 文件，將相關(guān)內(nèi)容粘貼至此。在這個過程中，所有包含“l(fā)lama”的模型都被替換為“qwen”。

2. 接下來，我們對項目中的 weight.py 的 load_from_hf_qwen 函數(shù)進(jìn)行修改，目的是逐步將 HuggingFace 的權(quán)重名稱與 TensorRT 的權(quán)重名稱對齊。執(zhí)行 build.py 后，雖然編譯成功，但執(zhí)行 run.py 的結(jié)果并未如預(yù)期。

3. 參照 TensorRT-LLM 的 docs/source/2023-05-19-how-to-debug.md 文檔，我們對模型進(jìn)行了詳細(xì)的調(diào)試，從外到內(nèi)打印模型層的數(shù)值，觀察 mean/sum/shape，并與原版進(jìn)行對比。經(jīng)過排查，我們發(fā)現(xiàn) attention 部分已經(jīng)包含了 rope 計算，通過調(diào)整 gpt attention plugin 的參數(shù)，最終使得輸出的 logits 正常。

4. 再次優(yōu)化 run.py，將 HuggingFace 原版 qwen_generation_utils.py 中的 make_context 函數(shù)遷移到 utils/utils.py 中，并導(dǎo)入該函數(shù)。這個函數(shù)被用來構(gòu)造一個 chat 版的 prompt 輸入，同時我們調(diào)整 eos 和 pad token 為 qwen 專屬的 <|im_end|> 或者 <|endoftext|>，最終 run.py 輸出也正常。

增加功能：Weight Only 量化

在 FP16 對齊成功，并且 run.py 以及 summarize.py 文件均能正常運(yùn)行之后，我們開始探索實現(xiàn) weight only int8/int4 量化。這只需要在 build.py 文件中對 weight only int8/int4 分支進(jìn)行輕微調(diào)整，包括 shape 的修改，以及保持權(quán)重名稱與 FP16 一致。接下來，我們進(jìn)行編譯測試，發(fā)現(xiàn)這一過程順利完成，且工作量并未超出預(yù)期，這部分工作基本無需投入大量人力資源。

增加功能：Smooth Quant

1. 在參考 Llama 項目的基礎(chǔ)上，我們將 hf_llama_convert.py 替換為 hf_qwen_convert.py 文件，該文件用于將 HuggingFace 的權(quán)重導(dǎo)出至 FasterTransformer (FT) 格式的文件。同時，我們將 llama/weight.py 中的 load_from_binary 函數(shù)重命名為 load_from_ft 復(fù)制到 qwen/weight.py 中，并根據(jù)我們導(dǎo)出的權(quán)重文件進(jìn)行了適當(dāng)?shù)男薷?。然后，我們?qwen/build.py 中默認(rèn)的加載函數(shù)從 load_from_hf_qwen 更改為 load_from_ft。為了保證兼容性，我們也對 load_from_ft 函數(shù)進(jìn)行了 fp16 以及 weight_only 的 int8/int4 的適配，其適配流程與之前的基本相同。當(dāng)開發(fā)者未導(dǎo)出 FT 權(quán)重時，系統(tǒng)會自動加載 load_from_hf_qwen 函數(shù)以生成 engine。

2. 在 smooth quant 的實現(xiàn)方面，我們參考了 example/llama 的 smooth quant 過程，同樣在 hf_qwen_convert.py 中添加了 --smoothquant 選項。通過調(diào)試 example/llama/hf_llama_convert.py 文件，我們觀察了 smooth_llama_model 函數(shù)的計算方法以及參數(shù)的 shape，發(fā)現(xiàn)其 mlp 的 gate 和 up 與 qwen mlp 的 w2/w1 layer 相對應(yīng)，并且 w1/w2 共享一個輸入。這部分的適配與之前的基本一致，唯一的區(qū)別是，attention 和 mlp 中需要量化的層需要進(jìn)行轉(zhuǎn)置，然后在 weight.py 的 load_from_ft 函數(shù)中再次轉(zhuǎn)置回來。

3. 我們再次分析了 example/llama 的 smooth quant 過程，并參考了其 build.py 文件，發(fā)現(xiàn)其中一個有一個 from tensorrt_llm.models import smooth_quantize 過程。在這個過程中，_smooth_quantize_llama 函數(shù)會替換掉 trt-llm 原本的模型結(jié)構(gòu)。因此，我們在 qwen/utils 目錄下建立了一個 quantization.py 文件，參考了 llama 的 SmoothQuantAttention，并復(fù)用了其 SmoothQuantRmsNorm，從而實現(xiàn)了 qwen 的 smooth quant 的全部過程。

優(yōu)化效果

精度

• 測試平臺：NVIDIA A10 Tensor Core GPU(24G 顯存) | TensorRT 9.0.0.1。

• TRT_LLM engine 編譯時最大輸入長度：2048， 最大新增長度：2048。

• HuggingFace 版 Qwen 采用默認(rèn)配置，未安裝，未啟用 FlashAttention 相關(guān)模塊。

• 測試時：beam=batch=1，max_new_tokens=100。

• 測試結(jié)果（該結(jié)果由 examples/qwen/summarize.py 生成。注：量化后分?jǐn)?shù)與原版分?jǐn)?shù)越接近，精度越好）：

性能

• 測試平臺：NVIDIA A10 Tensor Core GPU (24G 顯存) | TensorRT 9.0.0.1。

• 測試數(shù)據(jù)集為 ShareGPT_Vicuna_unfiltered，下載地址：https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json

• 生成速度（token/s）：此指標(biāo)不僅包括 generation 的過程，同時也計算了 context 階段時間，因此它表示的是每秒實際處理（理解輸入和生成輸出）的 token 數(shù)量。

• 吞吐速度（requests/s）：此指標(biāo)代表在極限情況下，無請求間隙時，系統(tǒng)平均每秒能處理的請求數(shù)量。

• 以下的測試包含多個 batch，主要用于測試特定顯卡的極限運(yùn)行情況，測試過程僅使用 TensorRT-LLM python 運(yùn)行時環(huán)境。

• HuggingFace 版 Qwen 采用默認(rèn)配置，未安裝，未啟用 FlashAttention 相關(guān)模塊。

• 當(dāng)最大輸入長度：2048， 最大新增長度：2048，num-prompts=100, beam=1, seed=0 時，BenchMark 結(jié)果如下：

圖 1：TensorRT-LLM 與 HuggingFace

吞吐以及生成對比

（吞吐加速比最高 4.25, 生成加速比最高 4.69）

• 當(dāng)最大輸入長度：1024， 最大新增長度：1024，num-prompts=100, beam=1, seed=0。BenchMark 結(jié)果如下：

圖 2：TensorRT-LLM 與 HuggingFace

吞吐以及生成對比

（吞吐加速比最高 4.57, 生成加速比最高 5.56）

總結(jié)

從整個開發(fā)過程的角度來看，NVIDIA TensorRT-LLM 已經(jīng)實現(xiàn)了相當(dāng)豐富的功能。它支持新模型的工作量不大，因為可以復(fù)用已有模型的相關(guān)代碼，只需要進(jìn)行少量的改動即可完成對新模型的支持。這表明了 TensorRT-LLM 具有很好的擴(kuò)展性。此外，在精度方面，它能夠與 HuggingFace 保持一致，但在速度方面最高可以達(dá)到 HuggingFace 的 5.56 倍。綜合考慮這些因素，可以說 TensorRT-LLM 完全有資格成為大規(guī)模語言模型推理框架的首選。它極大地緩解了推理和部署的難題，為廣泛應(yīng)用大語言模型提供了有力支持。

項目代碼

開源地址：

https://github.com/Tlntin/Qwen-7B-Chat-TensorRT-LLM

關(guān)于作者

?

鄧順子

廣州大學(xué)工程管理專業(yè)，擁有管理學(xué)學(xué)士學(xué)位，目前擔(dān)任 NLP 算法工程師，主要研究留學(xué)教育領(lǐng)域的信息抽取與智能對話。曾獲得第二十一屆中國計算語言學(xué)大會（CCL2022）航旅縱橫杯一等獎（子任務(wù)二）和三等獎（子任務(wù)一），也是熱門 Rust 開源項目 Pake 的主要貢獻(xiàn)者之一。

?

趙紅博

河南科技大學(xué)機(jī)械制造專業(yè)，擁有工學(xué)學(xué)士學(xué)位，目前在 Boss 直聘擔(dān)任高性能計算開發(fā)工程師，主要研究招聘領(lǐng)域模型的推理加速工作。

?

季光

NVIDIA DevTech 團(tuán)隊經(jīng)理，博士畢業(yè)于中科院計算所。擅長 GPU 加速的視頻處理以及性能優(yōu)化，以及深度學(xué)習(xí)模型的推理優(yōu)化，在 GPU 視頻編解碼以及 CUDA 編程與優(yōu)化方面積累了豐富的經(jīng)驗。

GTC 2024 將于 2024 年 3 月 18 至 21 日在美國加州圣何塞會議中心舉行，線上大會也將同期開放。點擊“閱讀原文”或掃描下方海報二維碼，立即注冊 GTC 大會。