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

速度和精度

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

在 Perf 內容之前，需重申在做 PTQ 量化時需對哪些 OP 進行量化。以經典的 Transform 結構為例，量化主要圍繞紅色、藍色和綠色框進行，涉及 4 種 GEMM 運算和 Multi-Head Attention 的量化。

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

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

第三步為評估，使用 MMLU 進行估計。

針對第三步，做精度評估時，如圖所示，第一個小紅框對 MMLU 78 個子數據集做了評估。因為篇幅較大，省略了中間的數據集，只展示其中的一部分。第一行代表了所做的量化方案。第一列是 baseline，GEMM 運算采用的是 FP16，在整個表中，我們對比了 Attention 以外的 4 種 GEMM 運算和對應的 kvcache 開啟 FP8 情況下的精度。

首先是 FP8、INT8 weightonly + FP16 kvcache 及最后一列對應的綠色框?？梢钥吹剑思?FP8 方案能夠保持精度比較好的量化方式，其他的比如 INT8_sq，或者是 INT8 weightonly + INT8 kvcache 并不能保持很好的精度。

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

我們再看看加速比，第一列對比了 FP8 和 FP16，性能提升 1.5~1.7 倍。另外兩種方式的加速比都比較不錯，但是仍然沒有 FP8 高。采用 INT8 sq 或者 INT8 sq + INT8 kvcache 對于精度的保持可能并不會太好。因此，我們優先推薦純 FP8 的方案。

這里還測試了 GEMM + kvcache+FMHA 方案。當對 FMHA 進行 FP8 GEMM 運算 enable 時，對比純 FP8 與 FP16 FMHA 和 FP8 FMHA 這兩種方式的精度，采用純 FP8 方案，當開啟 FMHA 時，它的精度保持也是比較高的。

GEMM +kvcache + FMHA 對應的性能：因為開啟 FMHA 的 FP8 僅是針對首 token 的優化，首 token 的計算一般情況是一個 computer bound 問題，結果如上圖所示。我們在某款 GPU 上測試了 Llama2 7B 模型，input sequence 越大，開啟 FMHA 的 FP8，帶來的加速比越來越明顯。

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

量化工具 AMMO/Modelopt

接下來介紹下量化工具 AMMO，它的最新的名字是 Modelopt。FP8 PTQ 量化的方式，可以總結為三個步驟：

第一步，Calibrate pytorch model

第二步，生成 model_config

第三步，生成 engine

其中：

第一步最重要的 API 是 Quantized API，通過 Quantized API 可以生成 Scaling 的計算過程。關于這個過程，我們可以傳入一個模型，設置量化的 config，比如設置成 FP8。最后，準備好需要的 calibrate 數據。

第二步主要是幫助我們生成一個 Json 文件和一組 weight 文件。Json 文件主要存儲模型結構或者元數據。在 weight 文件中，group 的大小主要受 Tensor Parallelism 和 Pipeline Parallelism 影響，weight 則主要用來存儲對應的參數。這步最重要的是 API，直接調用一個 API，就可以轉成 model config，方便 TensorRT 生成 engine 時使用。

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

如何 Debug？

在使用過程中，如果遇到問題，該如何 debug？具體的 debug 過程如下：

第一步，找到想要輸出的 tensor 做注冊，這里的注冊通過一個 API 就可以完成。

第二步，build engine。

第三步，直接打開 debug model 進行打印即可。

如上圖所示，展示了一個簡單的 debug 過程。

另外，debug 可能會遇到一些經驗性問題：

在 debug 過程中，可能發現 GEMM 的輸出不對。這時，我們可以檢查 weights 的通道是否保持一致。因為 Huggingface 下載的不同模型，通道保持可能不太一樣。

Attention 輸出不對時，可以查看 attention 使用的 plugin 的參數，設置的是否正確。

Deep Dive

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

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

在這個過程中，大模型通常會有 6 個部分用到 FP8。其中模塊 1，4，5，6 為矩陣乘，2 是 FMHA，主要是 context phase 中的 batch GEMM 會用到 FP8。3 是 MMHA 中的 kvcache 會用 FP8 來存儲，以節省顯存。

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

我們剛開始拿到的是 FP16 的矩陣乘，針對這個矩陣乘的 Input 和 Weight 插入 QDQ 節點。對于 Output，如果使用 FP8 的 kvcache，也需要在 QKV GEMM 后面插入 QDQ 節點。如果不做 FP8 的 kvcache，或者矩陣乘是 QKV 之外的矩陣乘，由于 GEMM 的輸出是 half 型數據，因此不需要插入 QDQ 節點。

當把 QDQ 節點都插好之后，類似 TensorRT 的流程做 calibration，使用量化校準數據集作為模型的輸入，對每一個 activation 的 A-max 值做統計。我們并不是直接把 FP16 的數據 cast 成 FP8，而是通過一個量化的過程來完成。這里借助 Modelopt 工具中的 QDQ 來計算量化參數，也叫 Scaling Factor。有了 Scaling Factor，可以把左側插完 QDQ 的計算過程轉換成右側的計算過程。

其中輸入部分還是一個浮點的輸入，Quantized 節點可以把輸入量化成 FP8，在量化的過程中會盡可能與其他算子融合以減少數據傳輸。另外，權重矩陣用 weight 跟 Quantized scaling factor 乘完之后，存成一個 FP8 的值在顯存中。當計算矩陣乘時，可以把 FP8 weight load 進來，再把量化之后的 input 用 FP8 的 tensor core 進行計算。這里 FP8 只有 tensor core 支持，CUDA core 是沒有 FP8 的。用 FP8 tensor core 計算完之后，再做一個反量化，得出 FP16 的值。當然，輸出值的類型是根據實際需要來配置的，也可以是其他的數據類型。

在國內能買到的支持 FP8 的 H20 GPU 中，INT8 和 FP8 的算力峰值都是一樣的，都是 296 tflops。但實測中，FP8 用 Plug-in 或者用 TensorRT 融合的 myelin graph 運算，都會發現 FP8 比 INT8 快。這是因為 FP8 的計算是根據 Hopper 硬件的一些特性來做的計算。但是 INT8 很多的計算沒有參考最新 Hopper 的架構。所以，軟件優先級的問題也導致 FP8 矩陣乘的運算比 INT8 要快。當后續軟件層面也會優化 INT8，這個 Gap 將不存在。

除了矩陣乘，Attention 部分也可以借助 FP8 做運算。主要有兩個：

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

對于 FMHA，為什么不用 INT8？這里我們做過相應的實驗，INT8 的 FMHA 在精度上比 FP8 有很大的下降。所以，INT8 由于精度問題用不了，而 FP8 的精度更魯棒。同時，也因為 FP8 在絕對值相對較小的情況下，打點比 INT8 的數據分布更密集。但當絕對值很大時，對于離群點部分，INT8 不區分離群點和非離群點的打點密集程度，而 FP8 在離群點的地方打點很疏，在非離群點打點很密集，所以 FP8 的精度更魯棒。

FP8 中的 Quantized 和 Dequantized，有一個 per tensor 量化參數就可以搞定。不需像 INT8 per token + per channel 這樣復雜，FP8 就可以保持精度，這也是用 FP8 顯而易見的好處。

Masked Multi-Head Attention：Generation phase 計算 Attention 模塊時，需要用融合的算子。因為 MMHA 的計算量比 FMHA 小很多，雖然也需要做 batched GEMM，batched GEMM 的 batch 維度都是 BS * HEAD_NUM，區別在于，context phase 的 GEMM 是 [length, head_size] * [head_size, length]，而 generation 的 GEMM 是 [1, head_size] * [head_size, length]，因此 batch GEMM 并不是計算密集型的計算過程，所以換 FP8 的收益不大，直接用浮點即可。但是加載 KV-cache 的模塊可以通過 FP8 量化來節省顯存。KV-cache 有 INT8 KV-cache，也有 FP8 KV-cache。相比 INT8，FP8 的精度更魯棒，在 Hopper 硬件架構下，FP8 KV-cache 轉出浮點的速度比 INT8 快。所以，FP8 KV-cache 的 MMHA 速度比 INT8 KV-cache 的 MMHA 要快。

借助 NVIDIA NCU 工具，對比在未打開 XQA 情況下的 MMHA。圖中藍色代表 FP8 KV-cache，綠色代表 INT8 KV-cache。可以看到，INT8 的 MMHA kernel 在 XU pipe 上的利用率非常高，也就是所有的 kernel 運算，都會卡在這個地方，產生較高的瓶頸。（這里的 XU 是做 INT8 數值轉換用到的一個 pipe。）

FP8 主要用 ALU 和 FMA，bound 情況好于 INT8。所以，FP8 KV-cache 在數值轉換的 bound 程度相比 INT8 KV-cache 輕，所以 FP8 KV-cache MMHA 好于 INT8 KV-cache MMHA。

以上就是 FP8 在模型中的應用場景、優勢以及使用原因的簡要總結和介紹。

作者：

郝尚榮 |NVIDIA 解決方案架構師

趙一嘉 |NVIDIA 解決方案架構師