本文聚焦 NVIDIA FP8 訓練與推理的實踐應用。
FP8 訓練利用 E5M2/E4M3 格式,具備與 FP16 相當的動态範圍,适用于反向傳播與前向傳播。FP8 訓練在相同加速平台上的峰值性能顯著超越 FP16/BF16,并且模型參數越大,訓練加速效果越好,且其與 16-bits 訓練在收斂性和下遊任務表現上無顯著差異。FP8 訓練通過 NVIDIA Transformer Engine 實作,僅需少量代碼改動,并且支援 FlashAttention、混合精度訓練遷移等。支援 FP8 的架構包括 NVIDIA Megatron-LM、NeMo、DeepSpeed、飛槳 PaddlePaddle、Colossal AI、HuggingFace 等。
FP8 推理通過 NVIDIA TensorRT-LLM 實作,權重輸入先轉換為 FP8,并融合操作以提高記憶體吞吐,但部分輸出仍需 FP16 進行 reduction。NVIDIA 技術團隊正研究直接 FP8 reduction 以實作端到端的加速優化。
FP8 基本原理、采用理由和收益
圖 1. 四種資料類型
首先詳解 FP8 的概念,圖 1 展示了 FP8、FP16、FP32 與 BF16 四種資料類型。業界曾長期依賴 FP16 與 FP32 訓練,直至 GPT 橫空出世,BF16 因能避免計算過程中的數值溢出問題而受到青睐。
近年來,NVIDIA 技術團隊在 FP8 領域持續投入,釋出了多篇論文,并在曆屆 GTC 大會也分享了 FP8 在計算機視覺 (CV)、自然語言處理 (NLP) 以及大模型訓練中的實際效果。
圖 2. E4M3 與 E5M2 兩種資料格式
圖 2 表格展示了 E4M3 與 E5M2 兩種資料格式。其中可以看到,FP8 精度的 E5M2 資料格式的數部分,與 FP16 的保持一緻。這意味着 FP8 精度的 E5M2 資料格式具備與 FP16 相當的動态範圍,是以該資料格式常被用在訓練的反向傳播階段。而 E4M3 是在前向傳播中采用的 FP8 格式。圖 2 詳盡展示了 FP8 格式下各類特殊數值的表示方式。
當我們考慮浮點數的資料精度會不會損失的時候,這個浮點數往往會落入圖 2 下半部分裡粉色的 subnormal 區間。圖 2 下半部分是以 FP32 舉例的,讀者可根據圖 2 表格看到 FP8 的 subnormal 區間,是以我們在訓練模型時可進行理論分析,探究數值精度是否影響模型效果。
表 1. 援引的測試資料[1]僅供技術參考和讨論
表 1 旨在闡述采用 FP8 的原因,以在 NVIDIA H100 Tensor Core GPU 上為例,機關是 TFLOPS,相較 FP16 和 BF16,FP8 的峰值性能能夠實作翻倍。并且此表展示的基準測試資料是在 2023 年采集的,目前性能提升更為顯著。
圖 3. 測試資料僅供技術參考和讨論
圖 3 左側圖表對比了不同參數規模的 GPT-3 模型在 H100 上做 FP8 訓練,以及在 NVIDIA A100 Tensor Core GPU 上做 FP16/BF16 訓練的吞吐加速比。這個加速效果随模型規模正向變化,比如參數規模為 5B 至 40B,它的加速效果約為 2 到 3 倍。
右側表格則進一步對比了不同參數規模的模型同在 H100 GPU 上,使用 FP8 訓練相對 BF16 的性能加速比。就 126M 至 175B 參數的模型而言,除了個别特殊任務外,FP8 訓練的加速效果同樣随模型規模增大而提升。換言之,模型規模越大,采用 FP8 訓練的收益越大。
圖 4. 援引的測試資料[2]僅供技術參考和讨論
圖 4 援引的是行業測試資料。左側圖表顯示的是對 GEMM 單一計算任務的加速對比。在 H100 GPU 上 FP8 訓練相對于 A100 GPU 上 BF16 訓練的峰值性能加速比約為 6 倍,而在 GEMM 任務測試中接近 5 倍。并且鑒于底層 CUDA 核心持續優化,未來性能将進一步提升。
右側表格則展示了在不同規模的 GPT 模型做 FP8 訓練的實際加速效果,模型參數規模分别為 1B、3B、7B 和 30B。該圖表分别對比了在 H100 GPU 與 A100 GPU 上做 BF16 和 FP8 訓練的加速效果。可以看到 BF16 訓練對 1B、3B 模型的加速比約為 2.2 倍,而 FP8 訓練的加速比分别達 2.7 倍、2.8 倍,對 7B、30B 模型加速比則達到 3 倍和 3.3 倍,說明 FP8 訓練的性能優化效果更加顯著。
FP8 的訓練性能和收斂性
圖 5. 測試資料僅供技術參考和讨論
圖 5 展示了 FP8 訓練的性能與收斂性。右圖顯示在不同規模的 GPT 模型上使用 BF16 與 FP8 進行訓練的 loss (損失值)曲線,并以困惑度 PPL(Perplexity) 為度量名額。同色曲線代表相同模型規模,實線代表 BF16,虛線為 FP8。觀察 PPL 曲線走勢,可見随着訓練程序,FP8 與 BF16 的曲線幾乎完全吻合,表明兩者收斂性并無顯著差異。
左側表格則彙總了曆屆 GTC 大會中分享的下遊任務資料,包括 PPL 名額及 FP8 與 16-bits 訓練的對比,涵蓋 NLP 模型和 CV 模型。結果顯示,使用 FP8 訓練的模型與 16-bits 訓練的模型在各項名額上的數值差異甚微,證明了 FP8 訓練能達到同等效果。
圖 6. 援引的測試資料[3]僅供技術參考和讨論
圖 6 展示了我們在本地測試的一個 1.3B 參數模型的實際訓練結果,共進行了約 2.5 萬步訓練。結果顯示,該模型的 loss 曲線與預期基本相符,僅有微小(零點零幾)的差異。
這裡列舉在 FP8 訓練中實際采用的配置。可以看到使用 FP8 訓練時對代碼的改動極少,隻需添加幾行代碼即可,後文将詳細解釋這些代碼的具體含義。
--fp8-hybrid \
--transformer-impl transformer_engine \
--fp8-amax-history-len 1024 \
--fp8-amax-compute-algo max
此外,我們在實際訓練中的常見問題解答如下:
- 目前廣泛采用 BF16 進行混合訓練,轉用 FP8 是否需要自行編譯 kernel 或進行複雜的資料類型轉換?答案是否,建議使用 NVIDIA Transformer Engine 預置的多種 FP8 kernel(Linear、MLP、LayerNorm 等基礎算子及基于這些算子的fused kernel),無需開發,直接調用即可。
- 如果沒使用 NVIDIA Megatron 或 DeepSpeed 架構,而是采用自定義架構,可以無縫使用 Transformer Engine 進行 FP8 訓練嗎?答案是可以。隻需在 PyTorch 上使用 Transformer Engine 提供的 fp8_autocast 包裝器(wrapper),即可在原生 PyTorch 環境中開展 FP8 訓練。此 wrapper 主要用于提供一系列 FP8-safe 的算子,自動将高精度的輸入資料轉換為 FP8,簡化了低精度訓練的實作過程。在上述過程中,需要對每個 tensor 更新其縮放因子 (scale),為此我們引入 amax(maximums of absolute value)的概念,fp8_autocast wrapper 會更新 amax 值。此外,根據 amax 值,該 wrapper 還會自動計算每個 tensor 的實際 scale 值。
- Transformer Engine 除提供 FP8 layer-wise 子產品和自動資料類型轉換外,還有什麼功能?答案是它還支援 FlashAttention 機制。這意味着 Transformer Engine 也能夠提升傳統 BF16、FP16 訓練的性能。
- 對于已使用 BF16 訓練的存量模型,能夠使用 FP8 做繼續訓練嗎?答案是可以。實踐證明,BF16 格式的 checkpoint 可以直接導入進行 FP8 繼續訓練;反之亦然,即在預訓練階段使用了 FP8,那麼在 SFT(supervised Fine-Tuning) 階段,出于對模型精度或資料健壯性的考慮,仍舊可以從 FP8 無縫切換到 BF16 做繼續訓練。Transformer Engine 全面支援此類精度遷移的操作。
圖 7. 解讀 FP8 訓練中新增的五行代碼
圖 7 旨在解讀前文提及的 FP8 訓練中新增的五行代碼,代碼的功能是用于計算目前 tensor 的 scale 值。我們采用名為 delayed scaling 政策,即目前 tensor 的 scale 值并非基于實時計算得出,而是依據其曆史資料,例如基于前幾個疊代周期的值計算得出。計算方法可選擇取 max 值,也可采用最近時間的值。
以該圖展示的 amax history 說明,針對目前 tensor,系統可存儲 1,024 個 amax 值,并從中選取最大值作為目前 tensor 的 amax 值。随後,根據一個簡化的 recipe 算法即可計算出 scale 值。
實際應用中,Hopper GPU 上 FP8 訓練相較于 BF16 的加速效果為 30%-40%,低于 FP8 在單一 GEMM 計算任務中理論可達的 5 倍加速比。為解釋此現象,本文借助圖 8 進行闡述。
使用 Transformer Engine 訓練 FP8 LLM
圖 8. FP8 訓練在 Transformer Engine 上的完整流程
圖 8 顯示了訓練中前向與反向計算的精度差異:紅線表示高精度(BF16、FP32),綠線為 FP8。在整個訓練期間,圖檔上半部分的權重(weight)及下半部分的梯度(gradient)始終以高精度存儲。僅在執行 linear 操作時,才對目前 tensor 進行資料格式轉換(cast),轉為 FP8 精度計算,但 linear 輸出仍為高精度。是以,後續 bias 計算等均在高精度上進行。
圖示表明,實際訓練中僅 GEMM 計算采用 FP8,其餘計算保持高精度。盡管業界存在對非線性操作也采用 FP8 計算和存儲的激進政策,并在部分下遊任務中表現良好,但主流方案依然遵循上述精細化的精度配置設定原則。
目前支援 FP8 訓練的分布式訓練架構與工具包括 NVIDIA Megatron-LM、NeMo 架構,DeepSpeed、飛槳 PaddlePaddle、Colossal AI、HuggingFace 等,也就是說這些架構均已內建了 Transformer Engine,可選用上述任一架構進行大模型 FP8 訓練。
圖 9. 不同資料精度 集合 Transformer Engine 的訓練測試結果對比
圖 9 總結了上述重點,通過對比三類測試情況:綠線代表僅使用 BF16 訓練,橘線表示 BF16 訓練結合 Transformer Engine(即在啟用 FlashAttention 的同時,使用 Transformer Engine 内置的 fused kernel),藍線為 FP8 訓練結合 Transformer Engine。
綠線顯示,僅用 BF16 訓練時,模型在單 GPU 卡上即遭遇記憶體不足(OOM),而在啟用 Transformer Engine 後,依舊采用 BF16,模型也能順利完成訓練。若進一步轉為 FP8,單次疊代時間可提升約 34.56%。
中間的圖表展示了各類測試的顯存占用情況。如前文所述,權重、梯度及優化器(optimizer)的資料均以高精度存儲,此外,FP8 訓練因需在 checkpoint 中儲存額外值,訓練時顯存占用比 FP16 略高約 5% 以内。須注意,推理階段的顯存占用與訓練階段是完全不同的。
圖 10. Llama2-7B 模型做 FP8/BF16 繼續訓練的 loss 曲線高度一緻
圖 10 展示了對 Llama2-7B 模型做 FP8 繼續訓練的效果。本測試并未進行長時間的訓練,目的是在為了提供概念驗證 (PoC, Proof of Concept)。圖中共有四條曲線:灰色曲線代表全程使用 BF16 訓練,其餘三條線分别表示以 BF16 進行預訓練,儲存 checkpoint 後,再分别以 BF16 與 FP8 繼續訓練。從繼續訓練的兩條曲線來看,loss 曲線高度一緻,且與灰色曲線的趨勢也保持一緻。
圖 11. Llama2-7B 模型 1.3 萬步内 全程 FP8/BF16 訓練的 loss 曲線基本一緻
圖 11 展示的是對 Llama2-7B 在 1.3 萬疊代步内做全程 FP8 訓練,可以看到它和全程 BF16 訓練的 loss 曲線也幾乎一緻。
FP8 推理流程
本章節分享使用 TensorRT-LLM 進行 FP8 推理。前文圖 8 展示的 FP8 訓練在 Transformer Engine 上的完整流程,而在進入推理階段,圖 8 下半部分如梯度等訓練特有部分可去除,僅保留上半部份即可。
訓練時為確定梯度計算準确,權重通常維持為高精度(如 BF16 或 FP32),這是由于訓練時需更新參數,而在推理時,權重已固定,故可在模型加載或預處理階段提前将權重轉換為 FP8,確定模型加載即為 FP8 格式。此外,推理階段應盡量進行操作融合,如将 LayerNorm 與後續資料格式轉換操作整合,確定 kernel 輸入輸出盡可能維持 FP8,進而能夠有效提升 GPU 記憶體吞吐。同樣,GeLU (Gaussian Error Linear Unit) 激活函數也要力求融合。
目前少量輸出仍會保持為 FP16,原因是 NVIDIA NCCL 僅支援高精度規約操作 (reduction),是以現在仍然需采用 FP16 進行 reduction,完成後再轉化為 FP8。
圖 12. FP8 推理流程
經過上述融合後,推理流程就簡化為圖 12 所示。綠線代表 FP8 的輸入輸出(I/O),紅線表示高精度 I/O。圖中可見,最前端的 LayerNorm 輸出與權重均為 FP8,矩陣輸出暫時保持 FP16,與前文描述一緻。并且經過測試驗證可得,雖然矩陣輸出精度對整體性能影響較小,但與輸入問題的規模相關;且因其計算密集特性,對輸出形态影響微弱。
在完成 MHA(Multi-Head Attention)後,需要将結果轉換為 FP8 以進行後續矩陣計算,Reduction 是以 FP16 執行後再轉換到 FP8 的。對于 MLP1 和 MLP2,兩者邏輯相似,但不同之處在于:MLP1 的輸出可保持在 FP8,因為它已經把 GeLU 加 Bias 等操作直接融合到 MLP1 的 kernel。
由此引發的關鍵問題是,能否将剩餘紅線(高精度 I/O)全部轉為綠線(FP8 I/O),實作進一步的加速優化?這正是 NVIDIA 持續進行的方向。以 reduction 為例,NVIDIA 正研究直接實作 FP8 reduction,盡管中間累加仍需高精度,但在資料傳輸階段可采用 FP8。與現有 reduction 不同的是,FP8 reduction 内部需引入反量化(de-quantization)與量化 (quantization)操作,故需定制開發 reduction kernel。
最佳實踐:使用 TensorRT-LLM 實作 FP8 推理
TensorRT-LLM 是基于 NVIDIA TensorRT 建構,其 FP8 能力也主要是通過 TensorRT 提供。自 TensorRT 9.0 版本起,官方就已經開始支援 FP8 推理。要在 TensorRT 中啟用 FP8 推理,需完成以下幾步:
- 設定 FP8 标志:通過調用 config.set_flag (trt.BuilderFlag.FP8) 在 TensorRT 配置中啟用 FP8 支援。類似 INT8、BF16、FP16,FP8 也是類似的啟用方式。
- 添加 GEMM 縮放因子(scale):主要針對輸入和權重,需在 weight.py (TensorRT-LLM 中的檔案)中額外加載這些縮放因子。這是 FP8 推理中不可或缺的步驟。
- 編寫 FP8模型:現階段我們需要明确編寫需要 FP8 支援的模型。具體做法如下:将原始 FP16 輸入量化至 FP8,随後進行反量化;權重同樣進行量化與反量化操作。如此編寫的模型,TensorRT 會自動将量化與反量化操作盡可能與前一個 kernel 融合,以及将反量化操作與 matmul kernel 融合。最終生成的計算圖表現為量化後的 X 與 W 直接進行 FP8 計算,輸出也為 FP8 結果。
為了簡化 FP8 在 TensorRT-LLM 中的應用,TensorRT-LLM 已對其進行封裝,提供了 FP8 linear 函數和 FP8 row linear 函數來實作。對于使用直接線性層(linear layer),則無需重新編寫代碼,直接調用函數即可。
圖 13. FP8 推理計算流程
本文用圖 13 總結上述内容。首先權重以 FP8 精度存儲的,在進行計算前,權重先經曆一次反量化。注意,在此之前,權重的量化已在輸入前完成了,此處僅需進行反量化操作。這意味着,在進行矩陣内部計算時,實際上是使用反量化後的資料,通常是 FP16 或甚至 FP32 來進行運算的。
矩陣層盡管以 FP8 表示,但累加是采用 FP32 完成,累加後再乘以 scale 的相關參數,形成如圖所示的計算流程。最終得到的結果具備較高精度。由于累加器(accumulator)需要采用高精度的數值,是以,要獲得最終 FP8 的輸出結果,模型還需經過一個量化節點 (quantitation node)。
回顧整個流程,輸入經曆了量化與反量化操作。其中,量化 kernel 發生在反量化 kernel 之前,而 TensorRT 則會智能地融合這些 kernel,確定計算的高效和準确。
使用 Tensor-LLM 實作 FP8 推理的性能
表 2 測試資料僅供技術參考和讨論
表 2 對比第一列不同的 batch size,其中 max 值指的是在設定輸入為 1,024,輸出為 256,模型為 GPT-J 6B,所能使用的最大 batch size。
清單顯示,FP16 的 max 值為 75,而 FP8 的 max 值則提升至 85。原因是 FP8 僅節省了權重部分的記憶體,部分 tensor 以及 KV cache 仍保持在 FP16。表格最後一列展示了使用 FP8 KV cache 的情況,此時能夠看到其 max 值相比 FP16 的 max 值超出 2 倍。
在性能方面,單純啟用 FP8 會由于 batch size 提升有限,以及 KV cache 的影響,導緻性能提升并不顯著。然而,一旦将 KV cache 也轉換至 FP8,通過減半其記憶體消耗,模型吞吐量可以相較 FP16 提升約兩倍左右,這是一個相當理想的性能提升幅度。
[1][2][3] Khudia, D. & Chiley, V. Benchmarking Large Language Models on NVIDIA H100 GPUs with CoreWeave (Part 1)[EB/OL]. Mosaic AI Research, 2023-04-27. [2024-04-23]. https://www.databricks.com/blog/coreweave-nvidia-h100-part-1.
關于作者
高慧怡
2020 年加入 NVIDIA Solutions Architect 團隊,從事深度學習應用在異構系統的加速工作,目前主要支援國内 CSP 客戶在大語言模型的訓練加速工作。
薛博陽
2019 年加入 NVIDIA DevTech 團隊,專注于語言模型的技術開發與優化。目前主要負責 TensorRT-LLM 的功能開發。