本文將分享 TensorRT-LLM 中低精度量化內(nèi)容，并從精度和速度角度對(duì)比 FP8 與 INT8。首先介紹性能，包括速度和精度。其次，介紹量化工具 NVIDIA TensorRT Model Optimizer（簡(jiǎn)稱 ModelOpt）及其快速實(shí)現(xiàn)量化功能的方法。第三部分為 Debug。隨后講解 FP8 的 Deep-Dive，并進(jìn)行總結(jié)。

速度和精度

在講解精度之前，先介紹 NVIDIA Hopper 架構(gòu)上的數(shù)據(jù)類型 FP8，它有兩種數(shù)據(jù)類型：E5M2 和 E4M3，在 TensorRT-LLM 中目前支持 E4M3。對(duì) Tensor Core 硬件來(lái)說(shuō)，相比于 FP32/FP16 作為輸入，F(xiàn)P8 在數(shù)據(jù)傳輸上具有優(yōu)勢(shì)。另外，GEMM 運(yùn)算可直接使用 8 比特，相對(duì)于 16 比特的 TF32 或者 FP16 更快，且支持更高累加精度以保證精度。

在 Perf 內(nèi)容之前，需重申在做 PTQ 量化時(shí)需對(duì)哪些 OP 進(jìn)行量化。以經(jīng)典的 Transform 結(jié)構(gòu)為例，量化主要圍繞紅色、藍(lán)色和綠色框進(jìn)行，涉及 4 種 GEMM 運(yùn)算和 Multi-Head Attention 的量化。

PTQ 量化需計(jì)算 Scaling Factor，Multi-Head Attention 中的 GEMM 運(yùn)算在 Scaling Facotr 為 1 就可以保持不錯(cuò)的精度（目前，TensorRT-LLM 中為了提高精度，在該部分做了 Scaling Factor 不為 1 的實(shí)現(xiàn)，本文內(nèi)容是以 FMHA 的 Scaling 為 1 的情況下的分析）。而藍(lán)色和紅色 GEMM 運(yùn)算需進(jìn)行 Scaling 計(jì)算。除此之外，我們要保存 kvcache，也可對(duì) kvcache 進(jìn)行 8bit 量化，但需進(jìn)行 Scaling 計(jì)算。

計(jì)算 Scaling Factor 的方法是使用 Quantize 腳本，添加如上圖所示兩個(gè)參數(shù)（--qformat fp8，--kv_cache_dtype fp8）即可進(jìn)行 FP8 Scaling 計(jì)算。對(duì)于 FMHA Attention 無(wú)需 Scaling，生成 Engine 使用“--use_fp8_context_fmha enable”即可快速生成 FP8、kvcache 和 GEMM 運(yùn)算功能。

第三步為評(píng)估，使用 MMLU 進(jìn)行估計(jì)。

針對(duì)第三步，做精度評(píng)估時(shí)，如圖所示，第一個(gè)小紅框?qū)?MMLU 78 個(gè)子數(shù)據(jù)集做了評(píng)估。因?yàn)槠^大，省略了中間的數(shù)據(jù)集，只展示其中的一部分。第一行代表了所做的量化方案。第一列是 baseline，GEMM 運(yùn)算采用的是 FP16，在整個(gè)表中，我們對(duì)比了 Attention 以外的 4 種 GEMM 運(yùn)算和對(duì)應(yīng)的 kvcache 開(kāi)啟 FP8 情況下的精度。

首先是 FP8、INT8 weightonly + FP16 kvcache 及最后一列對(duì)應(yīng)的綠色框。可以看到，除了純 FP8 方案能夠保持精度比較好的量化方式，其他的比如 INT8_sq，或者是 INT8 weightonly + INT8 kvcache 并不能保持很好的精度。

再看藍(lán)色框部分，對(duì)比純 FP16 和純 FP8 方案的精度情況，以及最后一行紅色框展示了平均的精度比較。

我們?cè)倏纯醇铀俦龋谝涣袑?duì)比了 FP8 和 FP16，性能提升 1.5~1.7 倍。另外兩種方式的加速比都比較不錯(cuò)，但是仍然沒(méi)有 FP8 高。采用 INT8 sq 或者 INT8 sq + INT8 kvcache 對(duì)于精度的保持可能并不會(huì)太好。因此，我們優(yōu)先推薦純 FP8 的方案。

這里還測(cè)試了 GEMM + kvcache+FMHA 方案。當(dāng)對(duì) FMHA 進(jìn)行 FP8 GEMM 運(yùn)算 enable 時(shí)，對(duì)比純 FP8 與 FP16 FMHA 和 FP8 FMHA 這兩種方式的精度，采用純 FP8 方案，當(dāng)開(kāi)啟 FMHA 時(shí)，它的精度保持也是比較高的。

GEMM +kvcache + FMHA 對(duì)應(yīng)的性能：因?yàn)殚_(kāi)啟 FMHA 的 FP8 僅是針對(duì)首 token 的優(yōu)化，首 token 的計(jì)算一般情況是一個(gè) computer bound 問(wèn)題，結(jié)果如上圖所示。我們?cè)谀晨?GPU 上測(cè)試了 Llama2 7B 模型，input sequence 越大，開(kāi)啟 FMHA 的 FP8，帶來(lái)的加速比越來(lái)越明顯。

再來(lái)看下做量化的耗時(shí)情況，我們?cè)?CNNDaily 數(shù)據(jù)集上做了測(cè)試。在這個(gè)數(shù)據(jù)集中，我們首先推薦用 512 的數(shù)據(jù)量，就可以很好的完成 FP8 保持精度的 calibration，其概耗時(shí)是 40 毫秒。這是在另一款 GPU 上做的測(cè)試，如果顯存比較大，我們可以讓 Batch size 變大一點(diǎn)，這時(shí) calibration 的時(shí)間可以變成秒級(jí)。

量化工具 AMMO/Modelopt

接下來(lái)介紹下量化工具 AMMO，它的最新的名字是 Modelopt。FP8 PTQ 量化的方式，可以總結(jié)為三個(gè)步驟：

第一步，Calibrate pytorch model

第二步，生成 model_config

第三步，生成 engine

其中：

第一步最重要的 API 是 Quantized API，通過(guò) Quantized API 可以生成 Scaling 的計(jì)算過(guò)程。關(guān)于這個(gè)過(guò)程，我們可以傳入一個(gè)模型，設(shè)置量化的 config，比如設(shè)置成 FP8。最后，準(zhǔn)備好需要的 calibrate 數(shù)據(jù)。

第二步主要是幫助我們生成一個(gè) Json 文件和一組 weight 文件。Json 文件主要存儲(chǔ)模型結(jié)構(gòu)或者元數(shù)據(jù)。在 weight 文件中，group 的大小主要受 Tensor Parallelism 和 Pipeline Parallelism 影響，weight 則主要用來(lái)存儲(chǔ)對(duì)應(yīng)的參數(shù)。這步最重要的是 API，直接調(diào)用一個(gè) API，就可以轉(zhuǎn)成 model config，方便 TensorRT 生成 engine 時(shí)使用。

第三步也是通過(guò)一個(gè) API 就可以完成，也就是加載上一步的 model config，直接生成 engine 結(jié)果。在這過(guò)程中，有一些隱藏的參數(shù)，比如訓(xùn)練的模型 TP/PP 比較大或者并行比較，在推理時(shí)，可以通過(guò) API 讓 TP/PP 變小。圖中是我們用 Modelopt 工具做 PTQ 量化時(shí)，一些簡(jiǎn)單的 API。

如何 Debug？

在使用過(guò)程中，如果遇到問(wèn)題，該如何 debug？具體的 debug 過(guò)程如下：

第一步，找到想要輸出的 tensor 做注冊(cè)，這里的注冊(cè)通過(guò)一個(gè) API 就可以完成。

第二步，build engine。

第三步，直接打開(kāi) debug model 進(jìn)行打印即可。

如上圖所示，展示了一個(gè)簡(jiǎn)單的 debug 過(guò)程。

另外，debug 可能會(huì)遇到一些經(jīng)驗(yàn)性問(wèn)題：

在 debug 過(guò)程中，可能發(fā)現(xiàn) GEMM 的輸出不對(duì)。這時(shí)，我們可以檢查 weights 的通道是否保持一致。因?yàn)?Huggingface 下載的不同模型，通道保持可能不太一樣。

Attention 輸出不對(duì)時(shí)，可以查看 attention 使用的 plugin 的參數(shù)，設(shè)置的是否正確。

Deep Dive

接下來(lái)，對(duì) FP8 的 workloads 進(jìn)行 deep dive，看模型什么地方用了 FP8，以及采用 FP8 之后的具體收益和為什么要這么用。最后介紹用 FP8 build 出的 engine 中 Scaling factor 和 tensor core 是怎么調(diào)用的。讓大家了解 FP8 的底層原理，進(jìn)而放心的去使用。

接下來(lái)介紹下從 FP16 模型 build FP8 Tensor-LLM engine 的過(guò)程。圖中黃色部分代表通過(guò) Modelopt toolkit 做 FP8 的權(quán)值轉(zhuǎn)換，存出 Model_config，再通過(guò) TensorRT-LLM 中的 From_json_file 和 Build_and_save 組件，將 Model_config 轉(zhuǎn)成 TensorRT-engine。

在這個(gè)過(guò)程中，大模型通常會(huì)有 6 個(gè)部分用到 FP8。其中模塊 1，4，5，6 為矩陣乘，2 是 FMHA，主要是 context phase 中的 batch GEMM 會(huì)用到 FP8。3 是 MMHA 中的 kvcache 會(huì)用 FP8 來(lái)存儲(chǔ)，以節(jié)省顯存。

上圖展示了從 FP16 矩陣乘變成 FP8 矩陣乘的過(guò)程：綠色代表 FP16 精度，黃色代表 FP8 精度，藍(lán)色是 FP32 精度，灰色代表融合的過(guò)程。

我們剛開(kāi)始拿到的是 FP16 的矩陣乘，針對(duì)這個(gè)矩陣乘的 Input 和 Weight 插入 QDQ 節(jié)點(diǎn)。對(duì)于 Output，如果使用 FP8 的 kvcache，也需要在 QKV GEMM 后面插入 QDQ 節(jié)點(diǎn)。如果不做 FP8 的 kvcache，或者矩陣乘是 QKV 之外的矩陣乘，由于 GEMM 的輸出是 half 型數(shù)據(jù)，因此不需要插入 QDQ 節(jié)點(diǎn)。

當(dāng)把 QDQ 節(jié)點(diǎn)都插好之后，類似 TensorRT 的流程做 calibration，使用量化校準(zhǔn)數(shù)據(jù)集作為模型的輸入，對(duì)每一個(gè) activation 的 A-max 值做統(tǒng)計(jì)。我們并不是直接把 FP16 的數(shù)據(jù) cast 成 FP8，而是通過(guò)一個(gè)量化的過(guò)程來(lái)完成。這里借助 Modelopt 工具中的 QDQ 來(lái)計(jì)算量化參數(shù)，也叫 Scaling Factor。有了 Scaling Factor，可以把左側(cè)插完 QDQ 的計(jì)算過(guò)程轉(zhuǎn)換成右側(cè)的計(jì)算過(guò)程。

其中輸入部分還是一個(gè)浮點(diǎn)的輸入，Quantized 節(jié)點(diǎn)可以把輸入量化成 FP8，在量化的過(guò)程中會(huì)盡可能與其他算子融合以減少數(shù)據(jù)傳輸。另外，權(quán)重矩陣用 weight 跟 Quantized scaling factor 乘完之后，存成一個(gè) FP8 的值在顯存中。當(dāng)計(jì)算矩陣乘時(shí)，可以把 FP8 weight load 進(jìn)來(lái)，再把量化之后的 input 用 FP8 的 tensor core 進(jìn)行計(jì)算。這里 FP8 只有 tensor core 支持，CUDA core 是沒(méi)有 FP8 的。用 FP8 tensor core 計(jì)算完之后，再做一個(gè)反量化，得出 FP16 的值。當(dāng)然，輸出值的類型是根據(jù)實(shí)際需要來(lái)配置的，也可以是其他的數(shù)據(jù)類型。

在國(guó)內(nèi)能買(mǎi)到的支持 FP8 的 H20 GPU 中，INT8 和 FP8 的算力峰值都是一樣的，都是 296 tflops。但實(shí)測(cè)中，F(xiàn)P8 用 Plug-in 或者用 TensorRT 融合的 myelin graph 運(yùn)算，都會(huì)發(fā)現(xiàn) FP8 比 INT8 快。這是因?yàn)?FP8 的計(jì)算是根據(jù) Hopper 硬件的一些特性來(lái)做的計(jì)算。但是 INT8 很多的計(jì)算沒(méi)有參考最新 Hopper 的架構(gòu)。所以，軟件優(yōu)先級(jí)的問(wèn)題也導(dǎo)致 FP8 矩陣乘的運(yùn)算比 INT8 要快。當(dāng)后續(xù)軟件層面也會(huì)優(yōu)化 INT8，這個(gè) Gap 將不存在。

除了矩陣乘，Attention 部分也可以借助 FP8 做運(yùn)算。主要有兩個(gè)：

Fused Multi-Head Attention：做 Context phase 時(shí)，Attention 計(jì)算中的 batch GEMM 可以用 FP8 計(jì)算。因?yàn)?FMHA 是一個(gè)融合的 kernel，由兩個(gè) batch GEMM 和中間的 softmax 組成。由于 softmax 是累加的過(guò)程，所以必須用高精度 FP32 處理。但對(duì)于 batch GEMM，可以直接借助 FP8 的 Tensor Core 計(jì)算，最終輸出是一個(gè) FP8 的輸出。這樣輸出的原因是 FMHA kernel 后，緊跟著一個(gè) FP8 的矩陣乘 project GEMM，可以直接接收 FP8 的輸出，所以直接輸出一個(gè) FP8 即可，減少了一次量化。

對(duì)于 FMHA，為什么不用 INT8？這里我們做過(guò)相應(yīng)的實(shí)驗(yàn)，INT8 的 FMHA 在精度上比 FP8 有很大的下降。所以，INT8 由于精度問(wèn)題用不了，而 FP8 的精度更魯棒。同時(shí)，也因?yàn)?FP8 在絕對(duì)值相對(duì)較小的情況下，打點(diǎn)比 INT8 的數(shù)據(jù)分布更密集。但當(dāng)絕對(duì)值很大時(shí)，對(duì)于離群點(diǎn)部分，INT8 不區(qū)分離群點(diǎn)和非離群點(diǎn)的打點(diǎn)密集程度，而 FP8 在離群點(diǎn)的地方打點(diǎn)很疏，在非離群點(diǎn)打點(diǎn)很密集，所以 FP8 的精度更魯棒。

FP8 中的 Quantized 和 Dequantized，有一個(gè) per tensor 量化參數(shù)就可以搞定。不需像 INT8 per token + per channel 這樣復(fù)雜，F(xiàn)P8 就可以保持精度，這也是用 FP8 顯而易見(jiàn)的好處。

Masked Multi-Head Attention：Generation phase 計(jì)算 Attention 模塊時(shí)，需要用融合的算子。因?yàn)?MMHA 的計(jì)算量比 FMHA 小很多，雖然也需要做 batched GEMM，batched GEMM 的 batch 維度都是 BS * HEAD_NUM，區(qū)別在于，context phase 的 GEMM 是 [length, head_size] * [head_size, length]，而 generation 的 GEMM 是 [1, head_size] * [head_size, length]，因此 batch GEMM 并不是計(jì)算密集型的計(jì)算過(guò)程，所以換 FP8 的收益不大，直接用浮點(diǎn)即可。但是加載 KV-cache 的模塊可以通過(guò) FP8 量化來(lái)節(jié)省顯存。KV-cache 有 INT8 KV-cache，也有 FP8 KV-cache。相比 INT8，F(xiàn)P8 的精度更魯棒，在 Hopper 硬件架構(gòu)下，F(xiàn)P8 KV-cache 轉(zhuǎn)出浮點(diǎn)的速度比 INT8 快。所以，F(xiàn)P8 KV-cache 的 MMHA 速度比 INT8 KV-cache 的 MMHA 要快。

借助 NVIDIA NCU 工具，對(duì)比在未打開(kāi) XQA 情況下的 MMHA。圖中藍(lán)色代表 FP8 KV-cache，綠色代表 INT8 KV-cache?？梢钥吹剑琁NT8 的 MMHA kernel 在 XU pipe 上的利用率非常高，也就是所有的 kernel 運(yùn)算，都會(huì)卡在這個(gè)地方，產(chǎn)生較高的瓶頸。（這里的 XU 是做 INT8 數(shù)值轉(zhuǎn)換用到的一個(gè) pipe。）

FP8 主要用 ALU 和 FMA，bound 情況好于 INT8。所以，F(xiàn)P8 KV-cache 在數(shù)值轉(zhuǎn)換的 bound 程度相比 INT8 KV-cache 輕，所以 FP8 KV-cache MMHA 好于 INT8 KV-cache MMHA。

以上就是 FP8 在模型中的應(yīng)用場(chǎng)景、優(yōu)勢(shì)以及使用原因的簡(jiǎn)要總結(jié)和介紹。

作者：

郝尚榮 |NVIDIA 解決方案架構(gòu)師

趙一嘉 |NVIDIA 解決方案架構(gòu)師