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

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

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

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

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

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

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

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

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

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

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

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

當(dāng)前 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 進(jìn)行大模型加速訓(xùn)練的過程。

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

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

與之前的一些更高精度的方法相比，比如 FP32、TF32、FP16、BF16 等格式，F(xiàn)P8 具有更高的 Flops 數(shù)值。理論估計(jì) FP8 相比 FP32 有四倍的算力提升，比 BF16 有兩倍的提升。在下面的表格中可以看到，在實(shí)際端到端訓(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 當(dāng)中。而 E4M3 是一種精度更高但動(dòng)態(tài)范圍較小的表達(dá)方式，因此它更適合在 forward 過程中處理 weights 和 activations。這種混合形式，可以在大模型的訓(xùn)練過程中根據(jù)情況靈活的運(yùn)用這兩種方式。對(duì)比以前進(jìn)行的混合精度或低精度訓(xùn)練，TF32 可以無縫替換 FP32，但到了 BF16 的 AMP 階段，我們不僅需要處理計(jì)算的低精度，還需對(duì)整個(gè) Loss 和梯度進(jìn)行 scaling。在 FP16 AMP 中，我們會(huì)針對(duì)整個(gè)網(wǎng)絡(luò)維護(hù)一個(gè) loss scale factor，而精度降至 8 比特時(shí)，就需要更精細(xì)地制定一套 recipe 來維護(hù) FP8 的精度表現(xiàn)，即在 FP8 訓(xùn)練過程中，我們需要進(jìn)行 per-tensor scaling。但是在進(jìn)行 per-tensor 時(shí)，會(huì)引入數(shù)值不穩(wěn)定的問題，因此我們需要謹(jǐn)慎處理。

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

我們也針對(duì)大模型訓(xùn)練的不同場(chǎng)景，對(duì) FP8 的收斂性進(jìn)行了測(cè)試和驗(yàn)證。

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

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

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

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

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

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

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

零一萬物的雙語 LLM 模型：

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

零一萬物是一家專注于大語言模型的獨(dú)角獸公司，他們一直致力于在 LLM 模型，及其基礎(chǔ)設(shè)施和應(yīng)用的創(chuàng)新。其可支持 200K 文本長度的開源雙語模型，在 HuggingFace 預(yù)訓(xùn)練榜單上，與同等規(guī)模的模型中對(duì)比表現(xiàn)出色[1]。在零一萬物發(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 倍的性能提升。

在推理方面，零一萬物基于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 倍的吞吐提升。

Inflection AI 的 FP8 訓(xùn)練

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

與同屬訓(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)練策略能夠保證模型正常收斂并取得良好的性能[2]。

此圖片由Inflection AI 制作，

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

結(jié)語

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