一、背景

之前我們已經總結過 FP8 的發展歷程，以及其在大規模語言模型（LLM）訓練和推理中的應用。如今，FP8 推理幾乎已經成為行業共識，許多 LLM 推理框架已經支持 FP8 推理，且多家公司也開源了相應的 FP8 模型。例如，Meta 在最近發布的 LLaMA-3 技術報告中就提到了 FP8 推理的應用。

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

與此同時，還比較少看到使用 FP8 進行 LLM 訓練的工作，目前看到的有如下幾個：

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

然而，也有一些工作對 FP8 訓練持懷疑態度，因此我們決定自己做一些實驗，以進一步驗證 FP8 訓練。此外，有關 FP8 訓練可能的問題，也可以參考 NVIDIA 的官方文檔：探索 FP8 訓練中 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 訓練的穩定性不如 BF16，而 FP8 的 Bit 數更少，可能導致更多穩定性問題。因此，作者認為降低精度的訓練方案必須具有與更高精度的訓練方案相似的訓練穩定性和超參敏感性，才能具有成本效益。同時，作者發現目前可用的 FP8 訓練方法不夠穩健，無法將它們用作當前方案的替代品。

PS：當然，作者也強調了使用 FP8 進行 LLM 推理是完全沒問題的。

2.2. 實驗

2.2.1 FP8 訓練實驗

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

2.2.2 降低 Bit 數實驗

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

如下圖 Figure 7 所示，進一步使用 E8M3、E8M4、E8M5 和 E8M6 訓練 LLaMA 7B，在 E8M5 和 E8M6 時才能保證相對的穩定性：

三、Megatron-LM FP8 訓練驗證

3.1. 摘要

對于上述論文中的實驗我們持懷疑態度，與我們之前看到的各種結論不符，因此決定進行相應的復現。同時我們也在考慮一個問題：如果無法充分保證 FP8 訓練的穩定性以及可比 BF16 訓練的精度，怎么權衡是否要使用 FP8 訓練。比如說，FP8 訓練相比 BF16 訓練可以加速 30%，但是 Loss 下降會慢一些，那么是否要使用 FP8 訓練呢？除此之外，我們也進一步測試了 GPT3 系列模型在不同 Batch Size 和 Seq Length 下 FP8 相比 BF16 訓練的加速比，以便為相關決策提供參考。

3.2. FP8 訓練 Loss 對比

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

如下圖所示為一個 1B 模型使用 FP8 訓練和 BF16 訓練的 loss 對比，總共訓練了 21K 個 Step，其 BF16 的 loss 基本上和 FP8 相當，并且收斂趨勢完全一致。當然，我們也發現 FP8 的 loss 始終會比 BF16 高一點：

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

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

3.3. FP8 訓練速度對比

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

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

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

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

• 紅色：表示加速比小于 1，通常是模型比較小，Batch Size 比較小的情況。
• 藍色：表示加速比大于 1 并且小于 1.3，通常是模型相當比較大或者 Batch Size 比較大。
• 綠色：表示加速比大于 1.3，同時是模型很大或者 Batch Size 很大。?

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

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

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

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

3.4. 零一萬物的實踐

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

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

四、參考鏈接

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

本文轉載自??AI閑談??，作者： AI閑談 ????