低精度訓練方法

Accelerate 提供了與 TransformersEngine、MS-AMP 和 torchao 包的整合，以便在指定的支援硬體上使用低精度方法進行訓練。本文件將指導你瞭解支援哪些硬體，如何配置你的 Accelerator 以利用低精度方法，以及你在訓練時可以期待什麼。

FP8 訓練的含義

要了解更多關於使用 PyTorch 和 Accelerate 進行 FP8 訓練的細節，請檢視關於為什麼這可能很困難的概念指南。但本質上，與在 BF16 中訓練不同，訓練模型的某些（或全部）方面可以使用 8 位而不是 16 位來執行。挑戰在於這樣做而不會降低最終效能。

這僅在特定的 NVIDIA 硬體上啟用，即

• 3000 系列消費級顯示卡之後的任何產品（例如 4090）
• 基於 Hopper 的 GPU 架構（例如 H100 和 H200）

這將導致使用的記憶體有所減少（因為我們已經將訓練某些部分所需的記憶體減少了一半），並且對於可以用支援 FP8 的層替換某些層的較大模型，吞吐量也應該有所增加。

配置 Accelerator

目前支援三種不同的 FP8 後端（TransformersEngine、torchao 和 MS-AMP），每種後端都有不同的功能和配置。

要使用其中任何一個，都使用相同的核心 API。只需將 mixed_precision="fp8" 傳遞給 Accelerator，或在 accelerate config 提示混合精度時輸入，或作為 config.yaml 檔案中 mixed_precision 鍵的一部分。

from accelerate import Accelerator
accelerator = Accelerator(mixed_precision="fp8")

預設情況下，如果你的環境中可用 MS-AMP，Accelerate 將自動使用它作為後端。要自己指定它（並自定義 FP8 混合精度設定的其他部分），你可以使用 RecipeKwargs 資料類之一，例如 utils.AORecipeKwargs、utils.TERecipeKwargs 或 utils.MSAMPRecipeKwargs；你也可以在你的配置 yaml 中/在 accelerate launch 期間明確指定它。

from accelerate import Accelerator
from accelerate.utils import MSAMPRecipeKwargs
kwargs = [MSAMPRecipeKwargs()]
# Or to specify the backend as `TransformersEngine` even if MS-AMP is installed
# kwargs = [TERecipeKwargs()]
# Or to use torchao
# kwargs = [AORecipeKwargs()]
accelerator = Accelerator(mixed_precision="fp8", kwarg_handlers=kwargs)

mixed_precision: fp8
fp8_config:
  amax_compute_algo: max
  amax_history_len: 1024
  backend: TE
  fp8_format: HYBRID
  interval: 1
  margin: 0
  override_linear_precision: (false, false, false)
  use_autocast_during_eval: false

配置 MS-AMP

在這兩者中，MS-AMP 通常更容易配置，因為它只有一個引數：最佳化級別。

目前，Accelerate 整合支援兩種最佳化級別，"O1" 和 "O2"（使用字母“o”，而不是零）。

• "O1" 將把權重梯度和 all_reduce 通訊轉換為 8 位進行，而其餘部分則以 16 位完成。這減少了一般的 GPU 記憶體使用並加快了通訊頻寬。
• "O2" 還會將一階最佳化器狀態轉換為 8 位，而二階狀態則為 FP16。（目前只支援 Adam 最佳化器）。這會盡力最小化最終的精度下降，並能節省最大的記憶體。

要指定最佳化級別，請透過設定 optimization_level 引數將其傳遞給 FP8KwargsHandler

from accelerate import Accelerator
from accelerate.utils import FP8RecipeKwargs
kwargs = [FP8RecipeKwargs(backend="msamp", optimization_level="O2")]
accelerator = Accelerator(mixed_precision="fp8", kwarg_handlers=kwargs)

或者在 accelerate launch 期間透過 --fp8_backend=msamp --fp8_opt_level=O2

同樣，這也可以在你的 config.yaml 中設定

mixed_precision: fp8
fp8_config:
    backend: MSAMP
    opt_level: O2

配置 TransformersEngine

TransformersEngine 有許多用於自定義 FP8 計算如何以及進行哪些計算的選項。支援的引數及其含義的完整列表可在 NVIDIA 的文件中找到，但為了方便起見，它們也作為 FP8KwargsHandler 的文件字串重新陳述。

Accelerate 試圖設定合理的預設值，但自己探索和調整各種引數可能會帶來更好的效能。

要使用它，請指定 backend="te" 並修改你想要的任何引數，作為 kwarg 處理器的一部分。

from accelerate import Accelerator
from accelerate.utils import FP8RecipeKwargs
kwargs = [FP8RecipeKwargs(backend="te", ...)]
accelerator = Accelerator(mixed_precision="fp8", kwarg_handlers=kwargs)

或者在 accelerate launch 期間透過 --fp8_backend=te ...。使用 accelerate launch --fp8_backend=te -h 檢視相關引數。

同樣，這也可以在你的 config.yaml 中設定

mixed_precision: fp8
fp8_config:
    amax_compute_algo: max
    amax_history_len: 1024
    backend: TE
    fp8_format: HYBRID
    interval: 1
    margin: 0
    override_linear_precision: (false, false, false)
    use_autocast_during_eval: false

配置 torchao

torchao 是一個由 PyTorch 驅動的、可修改的 FP8 後端，旨在比前兩個引擎更易於使用。與前兩個相比，ao 的一個核心區別是，為了數值穩定性，通常最好將模型的第一層和最後一層保持在常規精度（無論是 FP32 還是 BF16），然後將其他層量化到 FP8。因此，ao 的配置看起來有點不同。

注意：此 API 是實驗性的，可能會發生變化。

from accelerate import Accelerator
from accelerate.utils import AORecipeKwargs
kwargs = [AORecipeKwargs()]
accelerator = Accelerator(mixed_precision="fp8", kwarg_handlers=kwargs)

要了解有關要使用的具體引數的更多資訊，請參閱官方 torchao 倉庫。

示例集合

我們有示例展示了使用 accelerate 及其底層實現進行 FP8 訓練，這些示例可在 accelerate 倉庫中找到。目前我們支援的指令碼展示了

• 單 GPU
• 分散式資料並行（多 GPU）
• 完全分片資料並行
• DeepSpeed ZeRO 1 到 3

在此處瞭解更多資訊。

進一步閱讀

要了解有關在 FP8 中訓練的更多資訊，請檢視以下資源

• 我們的概念指南，詳細介紹了 TransformersEngine 和 MS-AMP
• transformers-engine 文件
• MS-AMP 文件
• torchao 文件