背景介紹

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

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

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

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

開(kāi)發(fā)與優(yōu)化過(guò)程

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

基礎(chǔ)功能支持

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

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

3. 參照 TensorRT-LLM 的 docs/source/2023-05-19-how-to-debug.md 文檔，我們對(duì)模型進(jìn)行了詳細(xì)的調(diào)試，從外到內(nèi)打印模型層的數(shù)值，觀察 mean/sum/shape，并與原版進(jìn)行對(duì)比。經(jīng)過(guò)排查，我們發(fā)現(xiàn) attention 部分已經(jīng)包含了 rope 計(jì)算，通過(guò)調(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ù)。這個(gè)函數(shù)被用來(lái)構(gòu)造一個(gè) chat 版的 prompt 輸入，同時(shí)我們調(diào)整 eos 和 pad token 為 qwen 專(zhuān)屬的 <|im_end|> 或者 <|endoftext|>，最終 run.py 輸出也正常。

增加功能：Weight Only 量化

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

增加功能：Smooth Quant

1. 在參考 Llama 項(xiàng)目的基礎(chǔ)上，我們將 hf_llama_convert.py 替換為 hf_qwen_convert.py 文件，該文件用于將 HuggingFace 的權(quán)重導(dǎo)出至 FasterTransformer (FT) 格式的文件。同時(shí)，我們將 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。為了保證兼容性，我們也對(duì) load_from_ft 函數(shù)進(jìn)行了 fp16 以及 weight_only 的 int8/int4 的適配，其適配流程與之前的基本相同。當(dāng)開(kāi)發(fā)者未導(dǎo)出 FT 權(quán)重時(shí)，系統(tǒng)會(huì)自動(dòng)加載 load_from_hf_qwen 函數(shù)以生成 engine。

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

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

優(yōu)化效果

精度

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

• TRT_LLM engine 編譯時(shí)最大輸入長(zhǎng)度：2048， 最大新增長(zhǎng)度：2048。

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

• 測(cè)試時(shí)：beam=batch=1，max_new_tokens=100。

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

性能

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

• 測(cè)試數(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 的過(guò)程，同時(shí)也計(jì)算了 context 階段時(shí)間，因此它表示的是每秒實(shí)際處理（理解輸入和生成輸出）的 token 數(shù)量。

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

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

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

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

圖 1：TensorRT-LLM 與 HuggingFace

吞吐以及生成對(duì)比

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

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

圖 2：TensorRT-LLM 與 HuggingFace

吞吐以及生成對(duì)比

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

總結(jié)

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

項(xiàng)目代碼

開(kāi)源地址：

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

關(guān)于作者

鄧順子

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

趙紅博

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

季光

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

GTC 2024 將于 2024 年 3 月 18 至 21 日在美國(guó)加州圣何塞會(huì)議中心舉行，線上大會(huì)也將同期開(kāi)放。點(diǎn)擊“閱讀原文”或掃描下方海報(bào)二維碼，立即注冊(cè) GTC 大會(huì)。