低精度訓練方法

新型硬體的釋出催生了能更好地利用這些硬體的新訓練正規化。目前，這體現為使用諸如 TransformersEngine (TE) 或 MS-AMP 等軟體包進行 8 位精度訓練。

為了介紹今天討論的主題，我們建議您回顧低精度使用指南，因為本文件將經常引用它。

快速圖表

以下是 MS-AMP 文件中的一個快速圖表，顯示了訓練期間每種解決方案的不同位精度。

最佳化級別	計算（GEMM）	通訊	重量	主權重	權重梯度	最佳化器狀態
FP16 AMP	FP16	FP32	FP32	不適用	FP32	FP32+FP32
Nvidia TE	FP8	FP32	FP32	不適用	FP32	FP32+FP32
MS-AMP O1	FP8	FP8	FP16	不適用	FP8	FP32+FP32
MS-AMP O2	FP8	FP8	FP16	不適用	FP8	FP8+FP16
MS-AMP O3	FP8	FP8	FP8	FP16	FP8	FP8+FP16

TransformersEngine 是第一個嘗試用 8 位浮點數進行訓練的解決方案。它的工作原理是為模型中的某些層使用可直接替換的層，利用其 FP8 引擎來減少位數（例如從 32 位減少到 8 位），而不會降低模型的最終精度。

具體來說，Accelerate 會查詢並用 TransformersEngine 版本替換以下層

結果，我們得到的模型中，大部分層是 BF16，而一些層是 FP8，從而減少了部分記憶體佔用。

根據經驗，我們注意到，在使用 TransformerEngine 時，只有當模型中絕大多數層都是由這兩個可替換的層構成時，效能提升才真正開始顯現。因此，只有當引數數量達到數十億級別或更高的大型模型才顯示出效能改進。

TransformerEngine 可以接收許多不同的引數，以自定義其執行 FP8 計算的方式和功能。下面是完整的引數列表：

margin：用於梯度縮放的邊距。
interval：用於重新計算縮放因子的頻率間隔。
fp8_format``：用於 FP8 recipe 的格式。必須是 HYBRID 或 E4M3 之一。（通常 HYBRID 用於訓練，E4M3` 用於評估）
amax_history_len：用於縮放因子計算的歷史記錄長度。
amax_compute_algo：用於縮放因子計算的演算法。必須是 max 或 most_recent 之一。
override_linear_precision：是否以更高精度執行 fprop、dgrad 和 wgrad GEMM 運算。

您可以將這些引數作為 utils.FP8RecipeKwargs 的一部分進行自定義，以幫助最佳化模型的效能。

如果我們注意前面提到的圖表，TE 只是將計算層轉換為 FP8，而其他所有部分都保持 FP32。因此，這最終會佔用最多的記憶體，但其好處是保證了訓練期間最終精度的損失最小。

MS-AMP 採用了與 TransformersEngine 不同的方法，它提供了三種不同的最佳化級別，以便將更多的操作轉換為 FP8 或 FP16。

基礎最佳化級別（O1）以 FP8 格式傳遞權重通訊（例如在 DDP 中），將模型權重儲存為 FP16，並將最佳化器狀態保留為 FP32。這個最佳化級別的主要好處是我們可以將通訊頻寬減少大約一半。此外，由於一半的資料被轉換為 FP8，權重被轉換為 FP16，從而節省了更多的 GPU 記憶體。值得注意的是，兩個最佳化器狀態都保持為 FP32。
第二個最佳化級別（O2）在此基礎上進一步改進，透過降低最佳化器狀態的精度。一個狀態是 FP8，另一個是 FP16。通常情況下，這已被證明只會在不降低最終精度的情況下帶來淨收益，同時提高了訓練速度，減少了記憶體佔用，因為現在每個狀態要麼是 FP16，要麼是 FP8。
最後，MS-AMP 還有一個第三最佳化級別（O3），它在 DDP 場景（如 DeepSpeed）中有所幫助。記憶體中的模型權重完全轉換為 FP8，主權重現在儲存為 FP16。這最大限度地減少了記憶體佔用，因為現在不僅幾乎所有東西都是 FP8，只剩下兩個狀態是 FP16。目前，只支援到 0.9.2 版本的 DeepSpeed，因此該功能未包含在 Accelerate 整合中。

還需要進行更多實驗，但有人指出，結合 MS-AMP 和 TransformersEngine 可能會透過依賴 NVIDIA 最佳化的 FP8 運算元並利用 MS-AMP 減少記憶體開銷的方式，實現最高的吞吐量。