FP8在NVIDIA GPU架構(gòu)和軟件系統(tǒng)中的應(yīng)用

在深度學(xué)習(xí)和人工智能的快速發(fā)展背景下，尤其是大語言模型（Large Language Model，LLM）的蓬勃發(fā)展，模型的大小和計(jì)算復(fù)雜性不斷增加，對(duì)硬件的性能和能效提出了極高要求。為了滿足這些需求，業(yè)界一直在尋求新的技術(shù)和方法來優(yōu)化計(jì)算過程。其中，F(xiàn)P8（8 位浮點(diǎn)數(shù)）技術(shù)憑借其獨(dú)特的優(yōu)勢(shì)，在 AI 計(jì)算領(lǐng)域嶄露頭角。本文作為 FP8 加速推理和訓(xùn)練系列的開篇，將深入探討 FP8 的技術(shù)優(yōu)勢(shì)，以及它在 NVIDIA 產(chǎn)品中的應(yīng)用，并通過客戶案例來展示 FP8 在實(shí)際部署中的強(qiáng)大潛力。

FP8 的原理與技術(shù)優(yōu)勢(shì)

FP8 是一種 8 位浮點(diǎn)數(shù)表示法，F(xiàn)P8 采取 E4M3 和 E5M2 兩種表示方式，其中 E 代表指數(shù)位（Exponent），M 代表尾數(shù)位（Mantissa）。在表示范圍內(nèi)，E4M3 更精準(zhǔn)，而 E5M2 有更寬的動(dòng)態(tài)范圍。與傳統(tǒng)的 FP16（16 位浮點(diǎn)數(shù)）和 FP32（32 位浮點(diǎn)數(shù)）相比，它顯著減少了存儲(chǔ)，提高了計(jì)算吞吐。

圖 1：浮點(diǎn)數(shù)據(jù)類型的結(jié)構(gòu)。所有顯示的值

（在 FP16、BF16、FP8 E4M3 和 FP8 E5M2 中）

都是最接近數(shù)值 0.3952 的表示形式。

圖 2：FP8 二進(jìn)制格式

數(shù)據(jù)表示位數(shù)的降低帶來了更大的吞吐和更高的計(jì)算性能，雖然精度有所降低，但是在 LLM 場(chǎng)景下，采用技術(shù)和工程手段，F(xiàn)P8 能夠提供與更高精度類型相媲美的結(jié)果，同時(shí)帶來顯著的性能提升和能效改善。

性能提升：由于 FP8 的數(shù)據(jù)寬度更小，減少了顯存占用，降低了通訊帶寬要求，提高了 GPU 內(nèi)存讀寫的吞吐效率。并且在相同的硬件條件下，支持 FP8 的 Tensor Core 可以在相同時(shí)間內(nèi)進(jìn)行更多次的浮點(diǎn)運(yùn)算，加快訓(xùn)練和推理的速度。

模型優(yōu)化：FP8 的使用促使模型在訓(xùn)練和推理過程中進(jìn)行量化，這有助于模型的優(yōu)化和壓縮，進(jìn)一步降低部署成本。

與 INT8 的數(shù)值表示相比較，F(xiàn)P8 在 LLM 的訓(xùn)練和推理更有優(yōu)勢(shì)。因?yàn)?INT8 在數(shù)值空間是均勻分布的，而 FP8 有更寬的動(dòng)態(tài)范圍， 更能精準(zhǔn)捕獲 LLM 中參數(shù)的數(shù)值分布，配合 NVIDIA Transformer Engine、NeMo 以及 Megatron Core 的訓(xùn)練平臺(tái)和 TensorRT-LLM 推理優(yōu)化方案，大幅提升了 LLM 的訓(xùn)練和推理的性能，降低了首 token 和整個(gè)生成響應(yīng)的時(shí)延。

FP8 在 NVIDIA GPU 架構(gòu)和軟件系統(tǒng)中的應(yīng)用

作為 AI 計(jì)算領(lǐng)域的領(lǐng)導(dǎo)者，NVIDIA 一直在推動(dòng)新技術(shù)的發(fā)展和應(yīng)用。FP8 技術(shù)在其產(chǎn)品中得到了廣泛的支持和應(yīng)用。

GPU 架構(gòu)與 Transformer Engine

NVIDIA GPU 無論是 Hopper 架構(gòu)、還是 Ada Lovelace 架構(gòu)都支持 Transformer Engine 進(jìn)行 FP8 的訓(xùn)練和推理。Transformer Engine 是一項(xiàng)專門為加速 Transformer 模型訓(xùn)練和推理而打造的軟件庫(kù)，應(yīng)用混合的 FP8 和 FP16/BF16 精度格式，大幅加速 Transformer 訓(xùn)練，同時(shí)保持準(zhǔn)確性。FP8 也可大幅提升大語言模型推理的速度，性能提升高達(dá) Ampere 架構(gòu)的 30 倍。

推理方面

TensorRT-LLM

TensorRT-LLM是 NVIDIA 針對(duì) LLM 推出的高性能推理解決方案。它融合了 FasterTransformer 的高性能與 TensorRT 的可擴(kuò)展性，提供了大量針對(duì) LLM 的新功能，而且能與 NVIDIA Triton 推理服務(wù)器緊密配合，使先進(jìn)的 LLM 發(fā)揮出更優(yōu)異的端到端性能。

TensorRT-LLM 可運(yùn)行經(jīng) FP8 量化的模型，支持 FP8 的權(quán)重和激活值。經(jīng) FP8 量化的模型能極好地保持原 FP16/BF16 模型的精度，并且在 TensorRT-LLM 上加速明顯，無論在精度上還是性能上，都明顯好于 INT8 SmoothQuant。

TensorRT-LLM 還為 KV 緩存（key-value cache）提供了 FP8 量化。KV 緩存默認(rèn)是 FP16 數(shù)據(jù)類型，在輸入文本或輸出文本較長(zhǎng)時(shí)，占據(jù)顯著的數(shù)據(jù)量；在使用 FP8 量化之后，其數(shù)據(jù)量縮小一半，明顯地節(jié)省顯存，從而可以運(yùn)行更大的批量。此外，切換到 FP8 KV 緩存可減輕全局顯存讀寫壓力，進(jìn)一步提高性能。

訓(xùn)練方面

NVIDIA NeMo

NVIDIA NeMo 是一款端到端云原生框架，用戶可以在公有云、私有云或本地，靈活地構(gòu)建、定制和部署生成式 AI 模型。它包含訓(xùn)練和推理框架、護(hù)欄工具包、數(shù)據(jù)管護(hù)工具和預(yù)訓(xùn)練模型，為企業(yè)快速采用生成式 AI 提供了一種既簡(jiǎn)單、又經(jīng)濟(jì)的方法。NeMo 適用于構(gòu)建企業(yè)就緒型 LLM 的全面解決方案。

在 NeMo 中，NVIDIA 集成了對(duì) FP8 的全面支持，使得開發(fā)者能夠輕松地在訓(xùn)練和推理過程中使用 FP8 格式。通過 NeMo 的自動(dòng)混合精度訓(xùn)練功能，開發(fā)者可以在保持模型準(zhǔn)確性的同時(shí)，顯著提高訓(xùn)練速度。此外，NeMo 還將陸續(xù)提供了一系列配套的工具、庫(kù)和微服務(wù)，幫助開發(fā)者優(yōu)化 FP8 推理的性能等。這些功能使得 NeMo 成為了一個(gè)強(qiáng)大而靈活的框架，能夠滿足各種 AI 應(yīng)用場(chǎng)景的需求。

Megatron Core

Megatron Core 是 NVIDIA 推出的一個(gè)支持 LLM 訓(xùn)練框架的底層加速庫(kù)。它包含了訓(xùn)練 LLM 所需的所有關(guān)鍵技術(shù)，例如：各類模型并行的支持、算子優(yōu)化、通信優(yōu)化、顯存優(yōu)化以及 FP8 低精度訓(xùn)練等。Megatron Core 不僅吸收了 Megatron-LM 的優(yōu)秀特性，還在代碼質(zhì)量、穩(wěn)定性、功能豐富性和測(cè)試覆蓋率上進(jìn)行了全面提升。更重要的是，Megatron Core 在設(shè)計(jì)上更加解耦和模塊化，不僅提高了 LLM 訓(xùn)練的效率和穩(wěn)定性，也為二次開發(fā)和探索新的 LLM 架構(gòu)提供了更大的靈活性。

Megatron Core 支持 FP8 的訓(xùn)練，很快也將支持 MoE 的 FP8 訓(xùn)練。

FP8 在 LLM 訓(xùn)練和推理的成功應(yīng)用

以下將詳細(xì)介紹幾個(gè)成功案例，并深入剖析其中的技術(shù)細(xì)節(jié)，以展現(xiàn) FP8 在推動(dòng) AI 模型高效訓(xùn)練和快速推理方面的巨大潛力。

FP8 在 LLM 的推理方面的成功應(yīng)用

Google 的 Gemma 模型與 TensorRT-LLM 完美結(jié)合

Google 一直在尋求提高其 LLM 的推理效率。為此，Google 與 NVIDIA 團(tuán)隊(duì)合作，將 TensorRT-LLM 應(yīng)用于 Gemma 模型，并結(jié)合 FP8 技術(shù)進(jìn)行了推理加速。

技術(shù)細(xì)節(jié)方面，雙方團(tuán)隊(duì)首先使用 NVIDIA AMMO 庫(kù)對(duì) Gemma 模型進(jìn)行了精細(xì)的量化處理，將模型的權(quán)重和激活值從高精度浮點(diǎn)數(shù)（FP16/BF16）轉(zhuǎn)換為低精度浮點(diǎn)數(shù)（FP8）。通過仔細(xì)調(diào)整量化參數(shù)，技術(shù)團(tuán)隊(duì)確保了 Gemma 模型在使用低精度進(jìn)行推理的時(shí)候不損失準(zhǔn)確性。這一步驟至關(guān)重要，因?yàn)樗苯佑绊懙侥Ｐ屯评淼男阅芎蜏?zhǔn)確性。

接下來，技術(shù)團(tuán)隊(duì)利用 TensorRT-LLM 對(duì)量化后的 Gemma 模型進(jìn)行了優(yōu)化。TensorRT-LLM 通過一系列優(yōu)化技術(shù)（如層融合、flash-attention）來提高模型的推理速度。在這個(gè)過程中，F(xiàn)P8 技術(shù)發(fā)揮了重要作用。由于 FP8 的數(shù)據(jù)寬度較小，它顯著減少了模型推理過程中所需的顯存帶寬和計(jì)算資源，從而提高了推理速度。同時(shí)，F(xiàn)P8 能夠降低存放激活值、KV 緩存所需要的顯存，增大我們推理時(shí)能夠使用的 batch size，進(jìn)一步的提升我們的吞吐量。

最終，通過 TensorRT-LLM 和 FP8 技術(shù)的結(jié)合，Google 成功在模型發(fā)布的同時(shí)實(shí)現(xiàn)了 Gemma 模型的高效推理。這不僅為用戶提供了更快的響應(yīng)時(shí)間和更流暢的語言處理體驗(yàn)，還展示了 FP8 在推動(dòng) AI 模型推理加速方面的巨大潛力。

圖 3：靜態(tài)批處理，延遲 vs 吞吐量

從上圖我們可以看到 FP8 相較于 FP16 帶來的收益。橫軸代表的是某個(gè)端到端的時(shí)間限制，縱軸代表的是吞吐量。對(duì)于生成式模型的性能評(píng)測(cè)，不能夠單純的只考慮吞吐量或是延遲。如果只追求吞吐量，導(dǎo)致延遲太長(zhǎng)，會(huì)讓使用者有不好的體驗(yàn)；相反的，如果只考慮延遲，則會(huì)導(dǎo)致沒辦法充分利用硬件資源，造成浪費(fèi)。

因此，常見的評(píng)測(cè)條件是，在一定的延遲限制下，達(dá)到最大的吞吐量。以上圖為例，如果將端到端的時(shí)間限制設(shè)為 15 秒，則只有 FP8 能夠滿足要求，這展示了在一些對(duì)延遲限制比較嚴(yán)格的場(chǎng)景下，F(xiàn)P8 加速帶來的重要性。而在 20 秒這個(gè)時(shí)間限制下，都使用 Hopper GPU 進(jìn)行推理時(shí)，F(xiàn)P8 對(duì)比 FP16 在吞吐量上能夠帶來 3 倍以上的收益。這是由于上面提到的，F(xiàn)P8 能夠在相同的時(shí)間限制下使用更大的 batch size，從而有更好的 GPU 利用率，達(dá)到更高的吞吐量。

圖 4：首 token 延遲 vs 基于每 dollar 的吞吐量

另外一方面，從性價(jià)比的角度來看，F(xiàn)P8 也帶來了很大的收益。上圖中，橫軸是產(chǎn)生第一個(gè) token 的時(shí)間限制，而縱軸則是每單位的資源能產(chǎn)生的 token 數(shù)量。請(qǐng)注意由于要處理長(zhǎng)輸入，產(chǎn)生第一個(gè) token 的時(shí)間通常會(huì)比較長(zhǎng)，而后面的階段可以透過流處理（streaming）的方式，每產(chǎn)生一個(gè) token 就立即返回給使用者，讓使用者不會(huì)覺得等待時(shí)間過長(zhǎng)。我們可以觀察到，隨著時(shí)間限制的增加，F(xiàn)P8 帶來的收益也越來越大。尤其是在 L4 Tensor Core GPU 上面，由于 GPU 本身的顯存較小，因此 FP8 節(jié)省顯存使用帶來的效益更加的明顯。

FP8 在 LLM 訓(xùn)練方面的成功應(yīng)用

Inflection AI 的 FP8 訓(xùn)練

Inflection AI 是一家專注于 AI 技術(shù)創(chuàng)新的公司，他們的使命是創(chuàng)造人人可用的 AI，所以他們深知大模型的訓(xùn)練對(duì)于 AI 生成內(nèi)容的精準(zhǔn)性和可控性至關(guān)重要。因此，在他們近期推出的 Inflection2 模型中，采用了 FP8 技術(shù)對(duì)其模型進(jìn)行訓(xùn)練優(yōu)化。

Inflection-2 采用了 FP8 混合精度在 5000 個(gè) NVIDIA Hopper 架構(gòu) GPU 上進(jìn)行了訓(xùn)練，累計(jì)浮點(diǎn)運(yùn)算次數(shù)高達(dá)約 10^25 FLOPs。與同屬訓(xùn)練計(jì)算類別的 Google 旗艦?zāi)Ｐ?PaLM 2 相比，在包括知名的 MMLU、TriviaQA、HellaSwag 以及 GSM8k 等多項(xiàng)標(biāo)準(zhǔn)人工智能性能基準(zhǔn)測(cè)試中，Inflection-2 展現(xiàn)出了卓越的性能，成功超越了 PaLM 2，彰顯了其在模型訓(xùn)練方面的領(lǐng)先性，同時(shí)也印證了 FP8 混合精度訓(xùn)練策略能夠保證模型正常收斂并取得良好的性能。

從 LLM 訓(xùn)練到推理， FP8 端到端的成功應(yīng)用

零一萬物的雙語 LLM 模型：FP8 端到端訓(xùn)練與推理的卓越表現(xiàn)

零一萬物是一家專注于 LLM 的獨(dú)角獸公司，他們一直致力于在 LLM 模型及其基礎(chǔ)設(shè)施和應(yīng)用的創(chuàng)新。其最新的支持 200K 文本長(zhǎng)度的開源雙語模型，在 HuggingFace 預(yù)訓(xùn)練榜單上，與同等規(guī)模的模型中對(duì)比表現(xiàn)出色。在零一萬物近期即將發(fā)布的千億模型 AI Infra 技術(shù)上，他們成功地在 NVIDIA GPU 上進(jìn)行了端到端 FP8 訓(xùn)練和推理，并完成了全鏈路的技術(shù)驗(yàn)證，取得了令人矚目的成果。

零一萬物的訓(xùn)練框架是基于 NVIDIA Megatron-LM 開發(fā)的 Y 訓(xùn)練框架, 其 FP8 訓(xùn)練基于 NVIDIA Transformer Engine。在此基礎(chǔ)上，零一萬物團(tuán)隊(duì)進(jìn)一步的設(shè)計(jì)了訓(xùn)練容錯(cuò)方案：由于沒有 BF16 的 baseline 來檢查千億模型 FP8 訓(xùn)練的 loss 下降是否正常，于是，每間隔一定的步數(shù)，同時(shí)使用 FP8 和 BF16 進(jìn)行訓(xùn)練，并根據(jù) BF16 和 FP8 訓(xùn)練的 loss diff 和評(píng)測(cè)指標(biāo)地差異，決定是否用 BF16 訓(xùn)練修正 FP8 訓(xùn)練。

由于 FP8 訓(xùn)練的過程中需要統(tǒng)計(jì)一定歷史窗口的量化信息，用于 BF16 到 FP8 的數(shù)據(jù)裁切轉(zhuǎn)換，因此在 BF16 訓(xùn)練過程中，也需要在 Transformer Engine 框架內(nèi)支持相同的統(tǒng)計(jì)量化信息的邏輯，保證 BF16 訓(xùn)練可以無縫切換到 FP8 訓(xùn)練，且不引入訓(xùn)練的效果波動(dòng)。在這個(gè)過程中，零一萬物基于 NVIDIA 軟硬結(jié)合的技術(shù)棧，在功能開發(fā)、調(diào)試和性能層面，與 NVIDIA 團(tuán)隊(duì)合作優(yōu)化，完成了在大模型的 FP8 訓(xùn)練和驗(yàn)證。其大模型的訓(xùn)練吞吐相對(duì) BF16 得到了 1.3 倍的性能提升。

圖 5：FP8 訓(xùn)練范式

在推理方面，零一萬物基于 NVIDIA TensorRT-LLM 開發(fā)了 T 推理框架。這個(gè)框架提供了從 Megatron 到 HuggingFace 模型的轉(zhuǎn)化，并且集成了 Transformer Engine 等功能，能夠支持 FP8 推理，大大減小了模型運(yùn)行時(shí)需要的顯存空間，提高了推理速度，從而方便社區(qū)的開發(fā)者來體驗(yàn)和開發(fā)。具體過程為：

將 Transformer Engine 層集成到 Hugging Face 模型定義中。

開發(fā)一個(gè)模型轉(zhuǎn)換器，將 Megatron 模型權(quán)重轉(zhuǎn)換為 HuggingFace 模型。

加載帶有校準(zhǔn)額外數(shù)據(jù)的 Huggingface 模型，并使用 FP8 精度進(jìn)行基準(zhǔn)測(cè)試。取代 BF16 張量以節(jié)省顯存占用，并在大批量推理中獲得 2~5 倍的吞吐提升。

圖 6：FP8 模型轉(zhuǎn)換與評(píng)測(cè)流程

總結(jié)與展望

基于上述 LLM 使用 FP8 訓(xùn)練和推理的成功實(shí)踐，我們可以看到 FP8 在推動(dòng) AI 模型的高效訓(xùn)練和快速推理方面的巨大潛力。隨著 FP8 訓(xùn)練和推理方法的不斷完善和廣泛應(yīng)用，我們確信 FP8 將在 LLM 應(yīng)用中扮演越來越重要的角色，敬請(qǐng)期待更多成功案例和技術(shù)解密。

審核編輯：劉清