一、背景

之前我們已經(jīng)總結(jié)過 FP8 的發(fā)展歷程，以及其在大規(guī)模語(yǔ)言模型（LLM）訓(xùn)練和推理中的應(yīng)用。如今，F(xiàn)P8 推理幾乎已經(jīng)成為行業(yè)共識(shí)，許多 LLM 推理框架已經(jīng)支持 FP8 推理，且多家公司也開源了相應(yīng)的 FP8 模型。例如，Meta 在最近發(fā)布的 LLaMA-3 技術(shù)報(bào)告中就提到了 FP8 推理的應(yīng)用。

FP8 推理相比 INT8 推理的最大優(yōu)勢(shì)在于其后訓(xùn)練量化（PTQ，Post Training Quantization）能夠獲得幾乎無損的精度，同時(shí)顯著提升推理速度。例如，相比 FP16，在 NVIDIA H20 上可以實(shí)現(xiàn) 2 倍的加速效果，而在 H100 上也可以獲得超過 1.5 倍的加速。

與此同時(shí)，還比較少看到使用 FP8 進(jìn)行 LLM 訓(xùn)練的工作，目前看到的有如下幾個(gè)：

• 微軟發(fā)布 FP8-LM 論文[2310.18313] FP8-LM: Training FP8 Large Language Models，并開源相應(yīng)的代碼 MS-AMP。
• 零一萬物在零一萬物面向萬卡集群的 AI Infra 建設(shè)中提到了 FP8 預(yù)訓(xùn)練。
• NVIDIA 和 Mistral AI 聯(lián)合發(fā)布 Mistral-NeMo-12B 模型，提到了使用 FP8 量化感知訓(xùn)練，但沒有介紹更多細(xì)節(jié)。
• NVIDIA 的 Megatron-LM 也早已通過 Transformer-Engine 庫(kù)支持了 FP8 訓(xùn)練。

然而，也有一些工作對(duì) FP8 訓(xùn)練持懷疑態(tài)度，因此我們決定自己做一些實(shí)驗(yàn)，以進(jìn)一步驗(yàn)證 FP8 訓(xùn)練。此外，有關(guān) FP8 訓(xùn)練可能的問題，也可以參考 NVIDIA 的官方文檔：探索 FP8 訓(xùn)練中 Debug 思路與技巧。

二、To FP8 and Back Again

2.1. 摘要

在 [2405.18710] To FP8 and Back Again: Quantifying the Effects of Reducing Precision on LLM Training Stability 中，作者指出，之前使用 FP16 訓(xùn)練的穩(wěn)定性不如 BF16，而 FP8 的 Bit 數(shù)更少，可能導(dǎo)致更多穩(wěn)定性問題。因此，作者認(rèn)為降低精度的訓(xùn)練方案必須具有與更高精度的訓(xùn)練方案相似的訓(xùn)練穩(wěn)定性和超參敏感性，才能具有成本效益。同時(shí)，作者發(fā)現(xiàn)目前可用的 FP8 訓(xùn)練方法不夠穩(wěn)健，無法將它們用作當(dāng)前方案的替代品。

PS：當(dāng)然，作者也強(qiáng)調(diào)了使用 FP8 進(jìn)行 LLM 推理是完全沒問題的。

2.2. 實(shí)驗(yàn)

2.2.1 FP8 訓(xùn)練實(shí)驗(yàn)

作者使用微軟開源的 https://github.com/Azure/MS-AMP.git（作者使用的是 v0.3.0，當(dāng)前最新的為 v0.4.0）來進(jìn)行 FP8 訓(xùn)練驗(yàn)證。如下圖 Figure 5 所示，作者使用 8 個(gè) H100 GPU 進(jìn)行實(shí)驗(yàn)，其中 MS-AMP 僅使用 O1 優(yōu)化，其在 GPT-2 124M 和 LLaMA 120M 上都有比較嚴(yán)重的收斂性問題，在 LLaMA 120M 上使用 FP8 訓(xùn)練甚至無法收斂

2.2.2 降低 Bit 數(shù)實(shí)驗(yàn)

如下圖 Figure 6 所示，使用 E8M3、E8M4 和 E8M5 來訓(xùn)練 TinyLLaMA 120M 模型，依然會(huì)出現(xiàn) Loss 不收斂的問題：

如下圖 Figure 7 所示，進(jìn)一步使用 E8M3、E8M4、E8M5 和 E8M6 訓(xùn)練 LLaMA 7B，在 E8M5 和 E8M6 時(shí)才能保證相對(duì)的穩(wěn)定性：

三、Megatron-LM FP8 訓(xùn)練驗(yàn)證

3.1. 摘要

對(duì)于上述論文中的實(shí)驗(yàn)我們持懷疑態(tài)度，與我們之前看到的各種結(jié)論不符，因此決定進(jìn)行相應(yīng)的復(fù)現(xiàn)。同時(shí)我們也在考慮一個(gè)問題：如果無法充分保證 FP8 訓(xùn)練的穩(wěn)定性以及可比 BF16 訓(xùn)練的精度，怎么權(quán)衡是否要使用 FP8 訓(xùn)練。比如說，F(xiàn)P8 訓(xùn)練相比 BF16 訓(xùn)練可以加速 30%，但是 Loss 下降會(huì)慢一些，那么是否要使用 FP8 訓(xùn)練呢？除此之外，我們也進(jìn)一步測(cè)試了 GPT3 系列模型在不同 Batch Size 和 Seq Length 下 FP8 相比 BF16 訓(xùn)練的加速比，以便為相關(guān)決策提供參考。

3.2. FP8 訓(xùn)練 Loss 對(duì)比

訓(xùn)練在 8*H100 機(jī)器進(jìn)行，訓(xùn)練數(shù)據(jù)集采用 cerebras/SlimPajama-627B · Datasets at Hugging Face。使用 NVIDIA 的 GitHub - NVIDIA/Megatron-LM: Ongoing research training transformer models at scale 訓(xùn)練框架，具體示例可以參考其 examples 中的 gpt3。

如下圖所示為一個(gè) 1B 模型使用 FP8 訓(xùn)練和 BF16 訓(xùn)練的 loss 對(duì)比，總共訓(xùn)練了 21K 個(gè) Step，其 BF16 的 loss 基本上和 FP8 相當(dāng)，并且收斂趨勢(shì)完全一致。當(dāng)然，我們也發(fā)現(xiàn) FP8 的 loss 始終會(huì)比 BF16 高一點(diǎn)：

如下圖所示，我們的結(jié)論與 Benchmarking Large Language Models on NVIDIA H100 GPUs with CoreWeave (Part 1) | Databricks Blog 中的結(jié)論基本一致：

除了 1B 模型之外，我們還驗(yàn)證了 345M 模型和 13B 模型，結(jié)論基本與上述一致。

3.3. FP8 訓(xùn)練速度對(duì)比

在驗(yàn)證收斂性的同時(shí)我們也同步驗(yàn)證了加速比，其 13B 模型 FP8 相比 BF16 可以獲得 30% 左右的加速，而 1B 模型可能只有 20% 左右，更小的模型加速比甚至小于 10%。（PS：不同的分布式策略都可能產(chǎn)生不同的結(jié)果，我們這里只是簡(jiǎn)單同配置下的驗(yàn)證）

如下圖所示，Benchmarking Large Language Models on NVIDIA H100 GPUs with CoreWeave (Part 1) | Databricks Blog 中 1B，3B 和 7B 模型的 FP8 訓(xùn)練相比 BF16 的訓(xùn)練加速比也只有 1.2x-1.3x：

為了對(duì)比不同配置下的性能，我們使用 Transformer-Engine 構(gòu)建了一個(gè) 1 層的 Transformer Block 進(jìn)行速度對(duì)比，同樣在 8*H100 上驗(yàn)證，采用 8TP，具體示例可以參考 Getting Started — Transformer Engine 1.8.0 documentation。

如下圖所示為 GPT-3 系列模型在 Seq Length=1024 時(shí)的性能，其中：

• 紅色：表示加速比小于 1，通常是模型比較小，Batch Size 比較小的情況。
• 藍(lán)色：表示加速比大于 1 并且小于 1.3，通常是模型相當(dāng)比較大或者 Batch Size 比較大。
• 綠色：表示加速比大于 1.3，同時(shí)是模型很大或者 Batch Size 很大。?

如下圖所示為 Seq Length 為 2048 的情況：

如下圖所示為 Seq Length 為 4096 的情況：

如下圖所示為 Seq Length 為 8192 的情況：

從上述結(jié)論可以看出，要想獲得比較大的加速比，通常需要具有比較大的模型或者比較大的 Batch Size、Seq Length。當(dāng)然，也并不是說 Seq Length 越大越好，可以看出，Seq Length 為 8K 是其加速比反而不如 4K。此外，也可以看出，大部分加速比不超過1.5x，甚至很多不超過 1.3x。（在實(shí)際使用中最好經(jīng)過一些充分的分析和實(shí)驗(yàn)）

3.4. 零一萬物的實(shí)踐

LLM 預(yù)訓(xùn)練的代價(jià)很高，比如可能需要上千個(gè) GPU 訓(xùn)練幾個(gè)月的時(shí)間，30% 的加速比似乎有很大的吸引力。然而，其結(jié)果又像薛定諤的貓，除非同時(shí)訓(xùn)練一個(gè) BF16 模型和 FP8 模型，才能確定 FP8 模型是否真的符合預(yù)期。

為了解決上述問題，零一萬物在 零一萬物面向萬卡集群的 AI Infra 建設(shè) 中提到了一個(gè) Trick 的方法。如下圖所示，每隔一段時(shí)間就會(huì) Load FP8 的 Checkpoint 并使用 BF16 進(jìn)行訓(xùn)練，驗(yàn)證 Loss 是否和 FP8 訓(xùn)練的 Loss 一致。如果出現(xiàn)不一致的情況，就會(huì)使用 BF16 的訓(xùn)練代替 FP8，并在一段時(shí)間后繼續(xù)使用 FP8 訓(xùn)練。最終作者獲得了 1.3x 的吞吐提升，不過并沒有說明這個(gè)提升是純粹的 FP8 相比 BF16 還是也包含了 BF16 的校驗(yàn)預(yù)算。

四、參考鏈接

1. https://arxiv.org/abs/2310.18313
2. https://01-ai.github.io/
3. https://mp.weixin.qq.com/s/ezdGxxmTRfEnzXmrVtwq7g
4. https://arxiv.org/abs/2405.18710
5. https://github.com/Azure/MS-AMP.git
6. https://huggingface.co/datasets/cerebras/SlimPajama-627B
7. https://github.com/NVIDIA/Megatron-LM
8. https://www.databricks.com/blog/coreweave-nvidia-h100-part-1
9. https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/examples/quickstart.html

本文轉(zhuǎn)載自??AI閑談??，作者： AI閑談 ????