本文主要介紹了 FP8 數據格式在大型模型訓練中的應用、挑戰及最佳實踐，展示了 FP8 在提升訓練速度和效率方面的潛力和實際效果。

FP8 格式

在介紹 FP8 格式之前，我們需要回答一個問題：為什么需要討論 FP8？從圖中可以看出，近年來大模型所需的算力急劇增長，從 GPT-1 到 GPT-3，再到類似 GPT-4 的 GPT MOE 1.8T，算力需求增長了數萬倍。這種增長速度的背后是硬件算力的提升。訓練過程中的一個重要指標是訓練時間。如果訓練一個模型需要半年甚至一年，這在實際操作中是不可行的，因為實際訓練時間可能是理論值的兩到三倍。因此，算力基礎設施的提升是大模型迅速發展的基礎。

從算力角度來看，近年來 GPU 的單卡算力提升了大約一千倍，這包括工藝制程的改進、硬件結構的優化以及更低的訓練精度。隨著 FP8 的引入，其 Tensor Core 算力是 FP16 的兩倍，為探索更大規模的模型提供了算力支持。

具體來說，FP8 的優勢包括：對于計算密集型算子，FP8 的 Tensor Core 相對于 BF16/FP16 能提供兩倍的算力，從而大大縮短計算時間；對于 Memory Bound 的算子，FP8 格式所需的數據量更少，可以節省緩存量，加快計算；如果將通信算子中的數據類型也替換成 FP8，也可以獲得一定的加速。最后，FP8 訓練的模型可以更好地與推理相結合，因為如果模型在訓練時的精度是 FP8，那么可以更快地部署到推理側，而不需要額外的 PTQ 量化過程。

FP8 數據格式包含兩種：E5M2 和 E4M3。E 代表指數位，M 代表尾數位。E5M2 包含五個指數位和兩個尾數位，而 E4M3 包含四個指數位和三個尾數位。E5M2 由于有更多的指數位，動態范圍更大；E4M3 有更多的尾數位，數值精度更好。這兩種數據格式在訓練時都有各自的應用場景。

在大模型訓練中使用 FP8

FP8 帶來了更快的訓練速度，但也對訓練精度提出了挑戰。接下來將介紹在大模型訓練中如何兼顧模型精度和訓練速度。

在介紹 FP8 之前，我們先回顧一下 16 位精度訓練中如何通過混合精度訓練來維持精度。這里列出了四種混合精度訓練的方法。第一種和最后一種嚴格來說不算混合精度，因為第一種是純 FP32 訓練，精度最好；最后一種是純 16 位精度訓練，速度最快。為了兼顧速度和精度，我們列出了額外的兩種模式：AMP-O1 和 AMP-O2。AMP-O1 相對于 O0 的不同點在于它會維護一份白名單，白名單中的 OP 會以低精度進行計算，如矩陣乘法和卷積算法，其他算子仍用高精度計算和存儲。AMP-O2 方案與 AMP-O3 更接近，不同點在于它會保留一些 unsafe 的 OP，這些 OP 會以 FP32 精度進行存儲和計算，如 LayerNorm 和 Softmax 等。此外，還會保留一份 FP32 類型的 Master Weight，因為在模型訓練后期，參數更新通常較慢，梯度值較小，容易出現大數加小數的問題，小數被吃掉，所以需要保留一份 FP32 的 Master Weight。目前 16 位精度訓練基本上都是采用 AMP-O2 的混合精度方法來訓練的。

FP8 訓練可以認為是一種 O1+O2 的混合模式。上邊這幅圖包含了前向和反向計算過程中的一些算子。紅色連接線表示高精度數據，綠色連接線表示低精度數據。無論是前向還是反向，整體訓練流程的精度仍然是 BF16 的，但會維護一份白名單，白名單中的 OP 以 FP8 精度計算，如 Linear Layer 中的矩陣乘法和 GELU。FP8 的 FMHA 目前在功能上是支持的，但在實際訓練過程中通常還是用高精度的 FMHA，以保證更好的收斂性。對于 FP8，我們看到它是以 O1 的模式嵌入 BF16 訓練的，BF16 訓練本身又是一個 O2 的混合精度方法，所以我們稱它為一種 O1+O2 的混合精度模式。

在訓練過程中，前向和反向采用了不同的數據精度。前向用 E4M3，因為前向時數值的動態范圍變化不大，用 E4M3 提供更好的精度；反向的梯度需要更大的動態范圍，所以反向用 E5M2 數據格式。這個流程圖中，藍色框表示從 BF16 到 FP8 的 Cast 過程。這個過程不像 FP32 到 BF16 那樣簡單直接。

接下來將詳細介紹 Cast 過程是如何實現的。因為 FP8 只有 4 位或 5 位指數位，小于 BF16，所以我們為了避免溢出的情況，在 Cast 過程中需要做量化。因為 FP8 的動態范圍有限，不足以表示模型的所有 Tensor，所以我們需要做 Per-tensor 的 Scaling。這與 FP16 不同，在 FP16 訓練中我們做的是全局的量化。圖中形象地表示了 Scaling 的過程。綠色中括號表示 E4M3 的動態范圍，能表示 2e-6 到 2e-8 次方范圍內的值。紫色中括號表示當前 Tensor 的數值分布。顯然，如果將 Tensor 從 BF16 直接轉換到 FP8，會有相當一部分值被直接 Flush 為 0，這些信息被丟棄，造成精度損失。我們的處理方式是給 Tensor 乘上一個系數，使 Tensor 的所有值向右平移，直到落到 E4M3 的表達范圍內。這樣，BF16 類型的 Tensor 就可以比較安全地 Cast 到 FP8。這個就是 Per-tensor Scaling 的過程，這個系數我們稱為 Scaling factor。

接下來面臨的問題是我們怎么來確定 Scaling factor。一種直接的方式是我們在計算得到一個高精度結果之后，通過類似于 torch.max() 這樣的一個算子，找到 Tensor 的最大值，通過這個最大值來計算 Scaling factor，然后再量化高精度的 Tensor 為 FP8 的輸出。但這種方法的問題是因為 Tensor 的 Shape 通常都會比較大，我們是沒有辦法把這個 Tensor 全部放到 GPU 的片上緩存 Shared Memory 中的。所以這個過程必須要借助 Global Memory 來進行數據的中轉，這就會帶來額外的一個訪存開銷。如果我們能提前知道 Scaling factor 的值，量化過程就可以提前到片上緩存 Shared Memory 中去完成。這時我們不需要等 Find Maximum 的值，Find Maximum 和 Scale 操作可以同時在片上緩存完成，從而避免額外的訪存開銷。

這里我們提前獲取 Scaling factor 的方式是 Delayed Scaling Recipe。這種方式的思想是通過當前 Tensor 的歷史迭代步信息來估計當前 Tensor 的最大值。

具體的，我們會建立一個 Amax History Buffer，記錄一個 Tensor 在歷史迭代步中的最大值。當需要當前 Tensor 的 Scaling factor 時，會從 History Buffer 中選出一個最大值，作為當前 Tensor 最大值的估計。有了最大值之后，可以計算 Scaling factor，從而對當前 Tensor 進行 FP8 量化。

另一方面，當要輸出 FP8 Tensor 時，我們會統計當前 Tensor 真實的最大值。將真實的最大值 New Amax 追加到 History Buffer 中。因為 History Buffer 是有長度的，所以當新的 Amax 追加到 History Buffer 末尾后，最前面的信息會被丟棄掉。這樣，可以一直用最近的歷史信息來估計當前 Tensor 的最大值。

接下來我們把 Delayed Scaling Recipe 過程放到一個真實的場景來介紹它是如何工作的：

圖中左邊部分是一個 Activation OP，輸入是一個高精度的 Activation，輸出是一個 FP8 的 Tensor。右邊是一個 Tensor Core 的 OP，輸入是 FP8 的 Tensor，輸出是一個高精度的值。這兩個 OP 可以類比到 Transformer Layer 里面的 Layer Norm 和 FC1。對于 Activation OP 來說，它的輸入和計算過程都是高精度的。當我們得到一個高精度的結果之后，我們會做兩件事：

第一件事，會統計當前 Tensor 的一個最大值，并將其追加到 History Buffer 中。同時另外一件事，我們會從 History Buffer 中選出一個最大值，作為當前 Tensor 最大值的估計，并計算出 Scaling factor，繼而將當前的 FP16 類型的 Tensor 量化到 FP8 進行輸出。對于 Tensor Core 的 OP 來說，它的 Activation 的輸入已經是 FP8 的 Tensor，權重也是用 Delayed Scaling Recipe 的方式來將其量化到 FP8。

這樣，我們將 GEMM 的所有的輸入都轉換成了 FP8，就可以用 FP8 的 Tensor Core 來進行計算，計算的結果是一個高精度的結果。

在輸出最終的結果之前，我們需要一個反量化的過程。這是因為 GEMM 的輸入對 Activation 和 Weight 都做了量化，所以它的值都被相應的左移或者右移。這就是 Megatron Core 框架里現在集成 FP8 訓練的一個方式。

FP8 訓練性能

接下來會介紹 FP8 訓練的性能結果：

這里使用的軟件鏡像是 NeMo Framework v24.01。我們可以看到在 Llama 模型上，FP8 訓練在訓練吞吐上的加速比在 33%-45% 范圍內。通過觀察 GPU 上的 Nsight System Report，發現 FP8 訓練的 Timeline 里面 Kernel 之間很容易出現氣泡。這個問題出現的原因是我們將代碼中最耗時的矩陣乘 Kernel 換成了 FP8，雖然它的計算時間減半了，但因為 FP8 的 Delayed Scaling Recipe 引入了一些和 Amax 以及 Scaling factor 相關的操作，引入了額外的 Kernel，所以導致 Host 端 launch Kernel 的 overhead 變大。此消彼長，使得 Kernel launch 跟不上 Kernel 計算的速度，產生 Launch bound 的問題。

如何解決 Launch bound 問題，我們將在后續的內容中介紹。

接下來是在另一張 GPU 上的測試結果，同樣也是在 Llama 模型上進行預訓練，鏡像是 NeMo Framework v24.01。可以觀察到 FP8 相對于 BF6 的加速比大約是 60%-73%。

最后是 MOE 模型上的一些 Benchmark 結果，模型是 Mixtral 8x7B，軟件是基于 Megatron-Core v0.7 開發的 FP8 版本。在這個版本上 ，FP8 的加速比達到了 63%。

再分享下性能上的最佳實踐：

FP8 這部分的性能問題并不多，比較常見的是 Kernel 之間的氣泡問題。為了解決這個問題，首先我們可以從減小 Host 端 Kernel launch overhead 的角度出發，盡可能地將這些 Kernel fuse 起來，減小 Kernel launch 的次數。比如我們可以將 Amax 以及 Scaling factor 相關的 Kerner fuse 起來，也可以把 Rotary Potential Embedding 這部分的 Kernel fuse 起來，以及 Swiglu 的 Fusion。除此之外，還可以利用 CUDA Graph 來將這些 Kernel 合并為一次 Graph 的 launch，來減小 Kernel launch 的開銷。另外，我們在代碼中要盡量避免 Host 端與 Device 端同步，這個同步會強制阻塞 Host 端操作，從而加重 Launch bound 問題。我們在平時寫代碼的過程中用 Torch 的 OP 可能就會不經意引入同步，我們在前向計算完成之后可能會對 Loss 進行一些處理，比如檢查一些 NaN 之類的，就會引入 Host 端與 Device 端同步。

最后是關于超參調整的建議：在顯存允許的情況下，我們會推薦嘗試用更多的 PP 而不是 TP，因為 PP 的通信粒度會更粗一些，引入的 Kernel 會更少。另一方面可以調整訓練的超參使得梯度累加的次數變少，當 GPU 的數量和問題規模比如 Global Batch Size 與 Synchronize 不變的情況下，一個 Global Step 的計算量是恒定的，當梯度累加次數越少時，意味著每次梯度累加所分到的計算量就越大。同時，每次梯度累加，Host 端 Kernel launch 的開銷是恒定的，所以當 Device 端的計算與 Host 端 launch 的開銷的比例達到一定程度，Launch bound 的問題就可以被減輕甚至直接消除掉。

這里關于超參調整的建議是從減小 Kernel 之間氣泡的角度出發的，實際在訓練過程中做超參調整時，要考慮的因素要更多，要做全盤的考慮。

FP8 訓練過程中的收斂性

接下來介紹 FP8 收斂性相關的信息，收斂性的結果我們按照大語言模型訓練的不同階段分別進行介紹：

首先是預訓練階段。我們訓練了一個模型結構和 Llama2-7B 相同的模型，數據集采用的是開源的 RedPajama，處理之后的數據量是 1.4T Token，超參用的是 TP2_PP1_DP128。

這里給出了從 300B 到 1.4T Token 之間的 Loss 曲線，以及訓練末期 FP8 和 BF16 Loss 的差值。

通過這兩幅圖可以看到 FP8 和 BF16 的 Loss 曲線是很接近的，二者之間 Loss 差值在10-3量級。

接下來是下游任務的結果，這里選取了 MMLU 和 LM-Harness 等幾個任務，對訓練好的模型進行評測。

結果顯示 FP8 和 BF16 訓練的模型在不同任務上的得分會有高低，但總體相差不大。所以，對于預訓練來說，FP8 訓練的模型無論是 Loss 曲線還是下游任務都可以和 BF16 匹配的很好。

接下來的場景是模型的增量訓練，比如從外部獲取到一個預訓練模型，我們需要給它注入新的知識，對 LLaMa 系列模型添加中文知識等等。我們選擇了從 HuggingFace 下載的 Llama2-7B 模型，并使用 Open-Web-Math 數據集進行訓練。訓練配置包括 TP1_PP2_DP4，我們提供了 FP8 和 BF16 的 Loss 曲線及其差值，可以看到同樣它們之間的 Loss 差值也非常小。

在下游任務評測中，我們選擇了 GSM8K 任務，并使用 OpenCompass 工具進行評測。我們每隔 500 步對訓練過程進行評分，以追蹤下游任務的變化。結果顯示，FP8 和 BF16 訓練的下游任務得分總體走勢一致，且都顯示出一定的波動性。兩次訓練中，FP8 和 BF16 的最高得分相近，表明 FP8 在增量訓練場景下的表現與 BF16 類似。

在 SFT（Supervised Fine-Tuning）結果中，我們選擇了 Llama2 系列的三個模型，并使用開源的三個數據集混合進行訓練。評測任務為 MT-Bench。從 Loss 曲線和下游任務結果來看，這三個模型的 FP8 Loss 曲線和下游任務得分都能與 BF16 對齊，證明了 FP8 在 SFT 訓練中的可行性。

在實際訓練過程中，我們沒有 BF16 baseline，因此需要通過其他方法判斷 FP8 是否處于正確的收斂路徑。一種方法是參考 01-AI 的做法，定期用 BF16 跑一定步數（如 100 到 200 步），作為 BF16 reference。通過比較 FP8 和 BF16 reference 的 Loss 曲線及下游任務得分，如果它們接近，則認為 FP8 訓練正確。如果差距較大，則用 BF16 替代 FP8 完成該期間的訓練。另一種方法是不一定需要 BF16 baseline，只要 FP8 訓練的 Loss 曲線持續下降，下游任務得分持續提升，就認為 FP8 處于正確的收斂路徑。

在對比 FP8 和 BF16 的下游任務得分時，應正確看待二者在下游任務上的得分差異。

Meta 最近發表的論文選取了上百個模型，這些模型除了初始化的隨機數種子不同外，其他配置和環境相同。論文統計了這些模型在不同下游任務 Benchmark 上的統計值，包括均值、方差、95% 的置信區間及單調性等。從方差及 95% 置信區間來看，不同隨機數種子對下游任務得分影響較大。例如，AGIEval 的平均得分是 23.44，95% 置信區間是 1.63，意味著模型得分在 21.8 到 25 之間都是合理的。GSM8K 的平均得分是 4.1，置信區間是 0.87，意味著模型得分在 3.2 到 4.9 之間都是合理的。盡管 FP8 和 BF16 之間的區別與隨機數種子的影響不同，但仍有助于設定 FP8 訓練的下游任務預期。當 FP8 的下游任務得分落在 BF16 的 95% 置信區間范圍內，應認定 FP8 訓練的模型與 BF16 匹配。

在收斂性過程中遇到問題的 Debugging Practices，可以將問題分為幾類。

首先是非 FP8 的問題，嘗試用 BF16 進行 Resume training，如果損失曲線與 FP8 一致，問題可能與 FP8 無關，而是與數據集或其他模塊相關。

第二類是軟件相關的 Bug，可以嘗試最新軟件棧或切換到 Transformer Engine/Megatron Core 的最近的幾個穩定版進行調試。

第三類是 Scaling factor 的問題，可以嘗試更保守的 Recipe，如 just-in-time 的 Scaling factor，如 Current scaling，來消除 Scaling factor 引起的誤差。另外，嘗試用 BF16 替代 FP8 來定位是哪個 GEMM 引起的問題。

最后是 Evaluation 過程中的問題，FP8 訓練的模型用 BF16 推理可能得到偏低分數。因為 FP8 數值格式的原因，一些精度會被丟棄，用 BF16 推理時這些信息會被重新引入，反而成為噪聲。對于精度要求高的任務，如 MMLU，會產生較大影響。推薦大家在用 BF16 推理遇到精度問題時不妨試試用 FP8 進行推理。如果訓練時的混合精度是 FP8 加 BF16，推理時則不能轉為 FP16 精度，因為 FP16 動態范圍有限，又沒有應用 Per-tensor scaling，容易出現上溢問題，所以需要推理精度仍然保持為 BF16。最后推薦用多點采樣的方式避免 Evaluation 過程中的噪聲問題。

展望

最后是展望和思考：

除了 Delayed Scaling 外，我們還進行了其他實驗，如 Current Scaling 和 Block Scaling。

Current Scaling 使用當前 Tensor 的最大值計算 Scaling factor，再將其 Cast 到 FP8，使用 Just-in-time 的 Scaling factor，對當前 Tensor 有更好的表示，不會出現上溢情況。

除此之外，我們還嘗試了更細粒度的 Scaling Recipe，如 Block Scaling 和 Per-channel Scaling，以每個 Block 或每行為一組，計算各自的 Amax 及 Scaling factor，保證不出現上溢情況下，減小下溢比例。因為 FP8 Tensor 附帶多組 Scaling factor，常規 FP8 GEMM 不支持這種情況，需要對 FP8 GEMM 進行改造，定制化開發來支持這種情況。

最后回顧下低精度訓練的發展過程與思考：

首先大語言模型經歷了從 32 位精度訓練到 16 位精度訓練的轉變，遇到了訓練不穩定性的問題。各個公司提出了不同解決方案，如 Google 通過跳過一定 Data batch 來消除 Loss spike，Meta 則通過修改 Learning Rate、Weight Decay 及模型結構等來優化。最終發現將數據類型從 FP16 換成 BF16 可以很好地解決這些問題，因為 BF16 具有更大的動態范圍，被廣泛應用于各大公司的大模型訓練上。

現在正在經歷 16 位精度到 8 比特精度轉化的過程，嘗試更穩定的 Scaling Recipe，如 Current Scaling 或 Block Scaling，或通過修改模型結構提高訓練穩定性。

總結

我們選擇更低精度的出發點是為了加快訓練速度，更快的訓練速度意味著可以用更多數據訓練更大模型，根據 Scaling Law 得到更好模型效果，或者在更短的時間內訓練出性能相當的模型。另一方面，低精度訓練格式天然對模型訓練效果有影響，因此需要找到方法使 FP8 訓練在絕大多數 Case 下穩定收斂，達到與高精度訓練相近的模型效果。

現在的 Delayed Scaling Recipe 在絕大多數場景下都可以很好的 Work，但仍有改進空間，無論是使用更魯棒的 Scaling Recipe，還是針對 FP8 訓練的特點調整模型結構，NVIDIA 技術團隊都在持續探索。無論如何，低精度訓練是大模型訓練的趨勢，NVIDIA 技術團隊將持續探索更好的 Scaling Recipe，讓大家更好地使用 FP8 訓練，相應進展會不定期的分享給大家。

作者簡介

劉宏斌

NVIDIA 加速計算專家，2020 年加入 NVIDIA DevTech 團隊，專注于 GPU 上深度學習模型的優化加速。目前主要負責生成式人工智能模型的訓練階段的加速優化。