如何使用FP8新技術(shù)加速大模型訓(xùn)練

利用 FP8 技術(shù)加速 LLM 推理和訓(xùn)練越來越受到關(guān)注，本文主要和大家介紹如何使用 FP8 這項新技術(shù)加速大模型的訓(xùn)練。

使用 FP8 進行大模型訓(xùn)練的優(yōu)勢

FP8 是一種 8 位浮點數(shù)表示法，F(xiàn)P8 的詳細介紹可以參考此鏈接：

https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/examples/fp8_primer.html#Introduction-to-FP8

其中，使用 FP8 進行大模型訓(xùn)練具有以下優(yōu)勢：

新一代 GPU 如NVIDIA Ada Lovelace、Hopper架構(gòu)配備了最新一代的 Tensor Core，可以支持 FP8 數(shù)據(jù)精度的矩陣運算加速。相比之前的 FP16 或 BF16 的數(shù)據(jù)類型，F(xiàn)P8 的 Tensor Core 可提供兩倍的 TFlops 算力。

除了計算上的性能加速之外，F(xiàn)P8 本身的數(shù)據(jù)類型占用的比特數(shù)比 16 比特或 32 比特更少，針對一些內(nèi)存占用比較大的 Operation，可以降低內(nèi)存占用消耗。

FP8 數(shù)據(jù)類型不僅適用于模型的訓(xùn)練，同樣也可用于推理加速，相對于以前常見的 INT8 的推理方法，使用 FP8 進行模型的訓(xùn)練和推理，可以保持訓(xùn)練和推理階段模型性能及數(shù)據(jù)算法的一致，帶來了更好的精度保持，避免了使用 INT8 進行額外的精度校正。

當然，F(xiàn)P8 對比 FP16 或者 FP32 在數(shù)值表示范圍上引入了新的挑戰(zhàn)，從上面的表格中可以看到，F(xiàn)P8 數(shù)據(jù)類型所能表示的數(shù)值范圍較小，精度較低。因此需要針對 FP8 引入更細粒度的算法改進，如針對每個 Tensor 進行 Scaling 的方法。對于 FP8 訓(xùn)練中的挑戰(zhàn)，NVIDIA 提出了一種 Delayed Scaling 的方法針對 FP8 Tensor 在訓(xùn)練過程中引入動態(tài) Scaling，使得在 FP8 訓(xùn)練過程中在加速矩陣運算的同時借助 per-Tensor scaling 的方法保持精度。

上述方法目前已被 NVIDIA 技術(shù)團隊實現(xiàn)，并集成到了Transformer Engine軟件包中。Transformer Engine 是 NVIDIA 提供的開源的訓(xùn)練工具包，專門針對 FP8 大模型訓(xùn)練實現(xiàn)了一系列功能，包含針對大模型所常見模型結(jié)構(gòu)如 Transformer 層等，同時針對 FP8 提供了 Delayed Scaling 這一方法的實現(xiàn)。

目前，Transformer Engine 已支持 PyTorch、JAX、Paddle 等主流框架，并與其它框架相兼容，且為了支持大模型訓(xùn)練，還實現(xiàn)了對模型及 Sequence Level 并行的方法。

使用 Transformer Engine 十分簡單方便，只需調(diào)用 Layer 層或 Transformer 層，并將 FP8 的 Delayed Scaling Recipe 包含在模型的定義的 context 中。剩下的訓(xùn)練過程中，所有 Tensor 的 Scaling 以及額外的輔助操作都可由 Transformer Engine 進行處理，無需額外操作 (參考上圖右側(cè)的示例)。

當前 Transformer Engine 已與NVIDIA NeMo、Megatron-LM以及HuggingFace 等業(yè)界開源社區(qū)訓(xùn)練框架融合，便于在大模型的訓(xùn)練中根據(jù)自己的需求方便調(diào)用 FP8 訓(xùn)練能力。比如：

在 NeMo 中想要打開 FP8 訓(xùn)練，只需要在配置文件中將 transformer_engine 和 FP8 分別設(shè)為 True，就可以方便的增加 FP8 的支持

在 Megatron-LM 中，只需要將 config 文件中的 FP8 設(shè)置為 hybrid，就可以用 FP8 進行大模型加速訓(xùn)練的過程。

FP8 旨在提升模型訓(xùn)練速度，目前已在 Hopper GPU 上對 Llama 系列模型進行 FP8 訓(xùn)練性能測評，結(jié)果顯示在 7B、13B 到 70B 等不同大小的模型下，使用 FP8 進行訓(xùn)練吞吐對比 BF16 其性能可提升 30% 至 50%。

FP8 在大模型訓(xùn)練中的特點，可簡單總結(jié)為以下幾點：

與之前的一些更高精度的方法相比，比如 FP32、TF32、FP16、BF16 等格式，F(xiàn)P8 具有更高的 Flops 數(shù)值。理論估計 FP8 相比 FP32 有四倍的算力提升，比 BF16 有兩倍的提升。在下面的表格中可以看到，在實際端到端訓(xùn)練任務(wù)的過程中，在不同的模型規(guī)模下，訓(xùn)練速度可以獲得約 1.37 倍到 1.52 倍的加速。

與更高精度的表示方法相比，F(xiàn)P8 有 E5M2 和 E4M3 兩種表示方式 (其中 E 為指數(shù)位，M 為尾數(shù)位)。E5M2 的指數(shù)位更多，意味著其數(shù)值表示范圍更大，梯度通常數(shù)值跨度更大，因此 E5M2 更適合用在 backward 當中。而 E4M3 是一種精度更高但動態(tài)范圍較小的表達方式，因此它更適合在 forward 過程中處理 weights 和 activations。這種混合形式，可以在大模型的訓(xùn)練過程中根據(jù)情況靈活的運用這兩種方式。對比以前進行的混合精度或低精度訓(xùn)練，TF32 可以無縫替換 FP32，但到了 BF16 的 AMP 階段，我們不僅需要處理計算的低精度，還需對整個 Loss 和梯度進行 scaling。在 FP16 AMP 中，我們會針對整個網(wǎng)絡(luò)維護一個 loss scale factor，而精度降至 8 比特時，就需要更精細地制定一套 recipe 來維護 FP8 的精度表現(xiàn)，即在 FP8 訓(xùn)練過程中，我們需要進行 per-tensor scaling。但是在進行 per-tensor 時，會引入數(shù)值不穩(wěn)定的問題，因此我們需要謹慎處理。

NVIDIA Transformer Engine 為用戶提供了相應(yīng)的 recipe，通過簡單傳入參數(shù)，即可方便地利用 FP8 的高算力，同時保持模型收斂性的表現(xiàn)。需要注意的是，并不是訓(xùn)練中的每個算子都要使用到 FP8，其主要應(yīng)用于線性層中的前向與后向矩陣乘運算中。而對于某些精度敏感的層，我們?nèi)詴褂酶呔扔嬎?，比如梯度更新、softmax 激活等。Transformer Engine 集成了很多 FP8 所需的可以保證精度的 recipe，并且 Transformer Engine 還集成到如 PyTorch、TensorFlow、Jax、Paddlepaddle 等更上層的訓(xùn)練框架，同時一些針對 LLM 訓(xùn)練的框架，如 Megatron-LM、NeMo Framework、DeepSpeed 等，也都集成了 FP8 能力。

我們也針對大模型訓(xùn)練的不同場景，對 FP8 的收斂性進行了測試和驗證。

上圖展示了一個從零開始預(yù)訓(xùn)練的損失曲線驗證，使用 Llama2 7B 模型，在 Pile 的 300 billion tokens 預(yù)訓(xùn)練數(shù)據(jù)集上，分別進行了 FP8 和 BF16 兩種精度下的模型訓(xùn)練，可以看到兩種精度的損失曲線吻合度極高，數(shù)值差異不到 1%。

此外，我們還使用 NVIDIA 開發(fā)的一個 8B 模型進行了繼續(xù)預(yù)訓(xùn)練測試，數(shù)據(jù)集同樣為 300 billion tokens，也可以看到 FP8 精度下和 BF16 的損失曲線差距也是很小的。同時在包括 MMLU 等多個下游任務(wù)上，也可以看到兩種精度所訓(xùn)練的模型的下游精度也是比較吻合的。

除了預(yù)訓(xùn)練階段，我們也對 SFT 階段的 FP8 訓(xùn)練精度進行了驗證，包括對 Llama2 7B、13B、70B 模型分別進行了 SFT (使用 NeMo 框架，數(shù)據(jù)集為開源社區(qū)中三個流行的英文數(shù)據(jù)集，MT-Bench 作為 SFT 精度驗證)。

可以看到對比了三種不同大小模型在兩種精度下的 SFT Loss 曲線，可以看到 Loss 曲線吻合度非常高，并隨著模型大小的增大，損失曲線明顯下降。

除了 Loss 曲線，也可以看到在 MT-Bench 測評集上三個模型在兩種精度下的 Score 也非常接近。

上圖是一個 SFT 模型生成效果的對比示例，可以看到在使用 13B 模型時，Prompt 為一個簡單編程任務(wù)的情況下，可以看到 FP8 和 BF16 生成的內(nèi)容也是非常接近和類似。

FP8 訓(xùn)練案例分享

零一萬物的雙語 LLM 模型：

FP8 端到端訓(xùn)練與推理的卓越表現(xiàn)

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

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

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

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

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

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

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

Inflection AI 的 FP8 訓(xùn)練

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

與同屬訓(xùn)練計算類別的 Google 旗艦?zāi)Ｐ?PaLM 2 相比，在包括知名的 MMLU、TriviaQA、HellaSwag 以及 GSM8k 等多項標準人工智能性能基準測試中，Inflection-2 展現(xiàn)出了卓越的性能，成功超越了 PaLM 2，彰顯了其在模型訓(xùn)練方面的領(lǐng)先性，同時也印證了 FP8 混合精度訓(xùn)練策略能夠保證模型正常收斂并取得良好的性能[2]。

此圖片由Inflection AI 制作，

如果您有任何疑問或需要使用此圖片，

結(jié)語

FP8 技術(shù)在推動 AI 模型的高效訓(xùn)練和快速推理方面有巨大的潛力，NVIDIA 的技術(shù)團隊也在和我們的客戶一起不斷探索完善應(yīng)用 FP8 訓(xùn)練和推理方法，未來我們也會持續(xù)為大家進行介紹以及最佳實踐分享。