混合精度訓練技術
1. 摘要與核心結論
深度學習正經歷從千萬引數到萬億引數的跨越,訓練算力需求以每 3.5 個月翻一番的速度急劇膨脹。在這一背景下,混合精度訓練(Mixed Precision Training)不再只是一個性能最佳化選項,而是已經成為大規模模型訓練不可或缺的工程基石。其核心思想看似簡單——在訓練過程中混合使用不同數值精度,如 FP16 與 FP32、BF16 與 FP32,乃至最新的 FP8——但在工程實現和數值理論層面蘊含著豐富的設計智慧。
本報告通過 15 個系統化章節,全面剖析混合精度訓練的技術原理、產業實踐和前沿演進。核心結論包括:第一,混合精度訓練可在幾乎不損失模型最終精度的前提下,將視訊記憶體佔用降至純 FP32 基線的約 50%,吞吐量提升 1.6× 至 2.15×。第二,該技術已成為支撐 GPT、Llama、Gemini 等千億乃至萬億引數模型高效訓練的效率基石,並被 PyTorch、TensorFlow、JAX 等主流架構深度整合。第三,損失縮放(Loss Scaling)是 FP16 混合精度的關鍵保護機制,而 BF16 則因其與 FP32 相同的指數範圍而天然規避了下溢風險,成為大語言模型訓練的首選格式。第四,隨著 Hopper 架構引入 FP8 支援,混合精度正進一步向 8 位浮點邁進,模型 Flops 利用率(MFU)可突破 60%~70%。第五,未來混合精度將由靜態策略向完全自適應的精度管理演進,並與稀疏化、量化感知訓練等技術深度融合,持續壓低大型模型訓練的邊際成本。
對於產業決策者而言,理解混合精度不僅是技術選型的需要,更是把握整個 AI 基礎設施效率脈搏的關鍵。本報告旨在提供一份兼具深度與廣度的技術參考,幫助讀者從硬體、架構、演算法三個維度建立完整的混合精度認知體系。
2. 技術背景與產業驅動:精度與效率的歷史博弈
深度學習訓練的數值精度選擇,本質上是一個在計算效率與數值穩定性之間尋找平衡的經典命題。要理解混合精度為何成為產業標配,必須回到 GPU 計算架構和模型規模增長的雙重歷史維度。
在 2017 年之前,深度學習訓練幾乎完全基於 IEEE 754 單精度浮點數(FP32)。FP32 擁有 1 位符號、8 位指數和 23 位尾數,可表示約 1.4×10⁻⁴⁵ 到 3.4×10³⁸ 範圍的數值,精度約為 7 個十進位制有效數字。對於當時數百萬引數的模型而言,FP32 提供了充分的數值精度,GPU 視訊記憶體也足以容納。然而,隨著 Transformer 架構的興起和模型規模的指數增長,情況發生了根本變化。
從 GPT-1 的 1.17 億引數到 GPT-4 的約 1.8 萬億引數,模型規模在五年間增長了四個數量級。與此同時,GPU 單卡視訊記憶體容量從 V100 的 32GB 增長至 H100 的 80GB,僅擴大 2.5 倍。這一剪刀差意味著,訓練大型模型必須採用模型並行、流水線並行等分散式策略,但即使如此,單裝置上的視訊記憶體壓力依然巨大。在訓練中,一個 FP32 張量佔用 4 位元組;對於萬億引數模型,僅引數本身就需要約 4TB 儲存,加上梯度、最佳化器狀態(如 Adam 的動量項和二階矩)和啟用值,總視訊記憶體需求輕鬆突破 20~30TB。若全部使用 FP32,需要上千張 GPU 才能容納,成本極為高昂。
更關鍵的是計算吞吐量的瓶頸。傳統 GPU 中,FP32 的算術邏輯單元(ALU)算力相對有限。以 V100 為例,其 FP32 峰值算力為 15.7 TFLOPS,而半精度 FP16 的峰值算力則為 125 TFLOPS,是前者的 8 倍。這一差異源於 NVIDIA 從 Volta 架構開始引入的 Tensor Core 專門單元,其原生支援 FP16 乘加運算,能夠在一個時鐘週期內完成 4×4 矩陣的乘加融合操作。如果訓練全程使用 FP32,意味著大量算力資源被閒置,訓練時間被不必要地拉長。
視訊記憶體頻寬是另一個關鍵約束。大型 Transformer 模型在訓練時,大量時間消耗在啟用值的讀寫上。將啟用值從 FP32 轉換為 FP16,資料量減半,頻寬壓力立即減輕,使得計算單元更接近峰值效能。這一效應在現代 GPU 上尤為顯著,因為計算能力的增長遠快於記憶體頻寬的增長。H100 的張量計算能力較 A100 提升了 3 倍以上,但 HBM3 頻寬僅從 2TB/s 提升至 3.35TB/s,這意味著若無精度最佳化,大量時間將浪費在等待資料上。
產業成本的壓力同樣不可忽視。訓練一個 GPT-4 級別模型,假設使用 25000 張 A100 訓練 90 天,電力成本可達數千萬美元。若能通過混合精度將訓練吞吐量提升 50%,即可節省超過 30% 的 GPU 租賃或折舊成本,價值數以千萬計。因此,從產業經濟角度,混合精度訓練帶來的效率提升是決定性的。
正是在這樣的技術推力與成本拉力的雙重作用下,混合精度訓練從 2017 年 Baidu 和 NVIDIA 聯合提出 FP16 混合精度訓練方法開始,迅速成為行業標準。2018 年 Google 在其 TPU v2/v3 中引入 BF16 支援,進一步解決了 FP16 的數值範圍不足問題。到 2023~2024 年,NVIDIA 的 Hopper 架構開始支援 FP8,混合精度的邊界被再次拓寬。
3. 混合精度訓練的基本範式:主權重高精化、前反向低精化、梯度尺度歸一化
混合精度訓練並非簡單地將所有張量轉換為低精度,而是遵循一個精妙的三層架構:主權重始終以高精度(FP32)儲存在視訊記憶體中,作為“真值”的唯一定錨;前向和反向傳播的計算和儲存則以低精度(FP16、BF16 或 FP8)進行,以充分利用硬體加速和頻寬節省;當梯度因低精度表示範圍不足而面臨下溢風險時,通過損失縮放將梯度拉回可表示區間,最終在最佳化器更新時還原為真實梯度。這一範式可概括為“主權重高精化、前反向低精化、梯度尺度歸一化”。
在一個典型的訓練迭代中,流程如下:
步驟一:主權重副本的維護。 在訓練的任意時刻,都有一份 FP32 的主權重儲存在視訊記憶體中。這是模型的“法律文本”,無論前向傳播用何種精度進行,所有引數更新的累積都作用在這份高精度副本上。這樣做確保了引數更新過程不會因半精度浮點數有限的尾數(FP16 僅 10 位有效尾數)而截斷微小更新。例如,當學習率為 1e-5 且梯度為 1e-8 時,兩數乘積為 1e-13,這在 FP16 中可能直接變為 0,而在 FP32 中則可被正常表示和累積。
步驟二:前向傳播的低精度化。 每個迭代開始時,從 FP32 主權重轉換(Cast)出一份 FP16/BF16 副本,輸入資料也通常轉換為同精度。所有的矩陣乘法、卷積等計算密集型操作,均使用低精度 Tensor Core 或等效硬體單元完成。由於半精度資料量僅為 FP32 的一半,矩陣乘法所需的視訊記憶體讀寫量大幅減少,同時 Tensor Core 的 FP16 算力通常是 FP32 的 8 倍或更高,因此前向計算時間顯著縮短。啟用值(即每層的輸出)通常以半精度儲存,僅在需要高精度的操作(如 Softmax 歸一化、LayerNorm)中臨時提升為 FP32,計算完成後再截斷回半精度儲存。
步驟三:損失計算。 損失函式通常涉及求和、對數等操作,數值範圍較寬。為保證訓練穩定性,損失值一般至少以 FP32 計算。實踐中常見兩種做法:一是將最後的 logits 輸出轉為 FP32 後計算損失;二是全程保持損失計算在 FP32 下進行。這樣能避免因 FP16 表示的精度侷限導致損失值發生偏移,進而影響收斂方向。
步驟四:反向傳播的低精度化。 基於低精度的啟用值和權重,反向傳播過程中產生的梯度張量也以半精度計算和儲存。這一步的視訊記憶體和頻寬節省同樣顯著。然而,半精度梯度可能因數值過小而無法表示,尤其是在訓練的後期或模型深處,梯度量級常常低於 FP16 的最小正規數(約 6×10⁻⁸)。若不加處理,這些梯度會被置為 0,造成引數停止更新的“死區”。
步驟五:損失縮放與梯度恢復。 為了解決上述下溢問題,在反向傳播前將一個損失縮放因子(Loss Scale)乘到損失值上,使得反向傳播產生的梯度被等比例放大。該因子通常為 8 到 65536 之間的 2 的冪次。反向傳播完成後,在將梯度傳遞給最佳化器之前,再將所有梯度除以同一縮放因子,恢復真實尺度。這一技巧在 FP16 混合精度中至關重要,而 BF16 由於其指數位數與 FP32 相同(8 位),動態範圍覆蓋了大部分梯度下溢場景,因此往往不需要複雜的損失縮放,甚至可以直接使用固定縮放因子 1.0。
步驟六:高精度引數更新。 還原後的梯度被轉換為 FP32,與 FP32 主權重一同送入最佳化器(如 AdamW)。最佳化器內部的狀態(動量、二階矩等)通常也保持 FP32,以保證更新計算的準確性。完成更新後,下一個迭代再次從 FP32 主權重轉換出半精度副本,迴圈往復。
這一基本範式的優雅之處在於,它將硬體的高效低精度計算與數值分析中的高精度累加原則有機結合,以最小的工程代價實現了效率與穩定性的兼得。這也是為什麼它能夠從學術界的初步探索,迅速演變為工業級大規模訓練的標準操作流程。
4. 主流數值格式對比:FP16、BF16 與 FP8 的權衡
混合精度訓練中“半精度”的具體選擇直接影響數值穩定性和硬體效率。當前產業實踐中,三種主要的浮點格式構成了遞進式技術路線:FP16、BF16(Brain Floating Point)和方興未艾的 FP8。
FP16 遵循 IEEE 754-2008 半精度標準,包含 1 位符號、5 位指數、10 位尾數。其表示範圍約為 5.96×10⁻⁸ 到 65504,精度約 3.3 個十進位制有效數字。FP16 的優勢在於硬體支援成熟,NVIDIA 從 Volta 架構開始就在 Tensor Core 中對其提供了原生加速,FP16 的矩陣乘加吞吐量可達 FP32 的 8 倍以上。然而,FP16 的指數位較少,表示範圍狹窄,使得它在處理大值(如上溢)和小梯度(下溢)時都面臨挑戰。這是引入損失縮放機制的直接原因。
BF16 由 Google 在 TPU 開發中提出,並迅速被 NVIDIA A100 及後續 GPU 採納。BF16 將尾數截斷為 7 位,但保留與 FP32 相同的 8 位指數,因此其表示範圍與 FP32 幾乎一致(約為 1.17×10⁻³⁸ 到 3.39×10³⁸),只是精度下降到約 2 個十進位制有效數字。這種設計使得 BF16 對梯度的下溢和上溢具有天然魯棒性,在大語言模型訓練中尤其受歡迎,因為 LLM 的梯度分佈往往比計算機視覺模型更偏向小值,且訓練過程中的超引數不需要因精度變化而大改。從 FP32 向 BF16 的遷移通常只需簡單地將張量截斷,無需複雜的動態損失縮放,訓練超引數幾乎可以保持不變。這大幅降低了使用者的使用門檻。但代價是,BF16 的尾數精度低於 FP16,在個別對精度要求極高的計算(如小學習率下的長尾收斂)中可能引入微小的收斂減速。
FP8 是 NVIDIA 在 H100 GPU 中引入的最新精度格式,分為 E4M3(4 位指數,3 位尾數)和 E5M2(5 位指數,2 位尾數)兩種變體。E4M3 提供更高的精度但範圍較窄,適用於前向傳播的權重和啟用;E5M2 提供更廣的範圍但精度低,更適合梯度表示。FP8 將資料寬度壓縮到 1 位元組,理論上相比於 FP16/BF16 可將儲存和頻寬需求再次減半,算力密度翻番。但 FP8 的範圍和精度更為逼仄,需要更細粒度的縮放策略來維持訓練穩定性。NVIDIA 為此在 Hopper 架構中引入了可程式設計的縮放因子,允許對輸入張量的每個 128 元素塊應用獨立的縮放,這種“塊級縮放”是 FP8 訓練成功的關鍵。目前,FP8 混合精度訓練已在 Llama 3 等前沿模型的訓練中得到初步應用,並展現出將 MFU 推至 60%~70% 的驚人效率。
三種格式並非互相取代,而是構成了工具箱,不同場景有不同最優選擇。對於 CV 和傳統中小模型,FP16 憑藉成熟的生態和充足的表示範圍依然廣泛使用;對於大語言模型和擴散模型,BF16 因其“零成本遷移”和穩定性成為預設選擇;對於追求極致效率和最新硬體的團隊,FP8 正成為新的效率高地。部分先進訓練方案甚至在前向傳播中使用 FP8,梯度使用 FP8,主權重使用 FP16,最佳化器狀態使用 FP32,形成多層級的混合精度策略,最大化硬體利用率。
5. 硬體加速生態:從 Volta 到 Hopper 的 Tensor Core 演進
混合精度訓練的產業落地與 GPU 硬體中專用低精度計算單元的演進密不可分。以 NVIDIA 為例,從 2017 年的 Volta 架構到 2023 年的 Hopper 架構,Tensor Core 已經歷了四代迭代,每一代都深刻塑造了混合精度訓練的能力邊界。
Volta 架構(V100,2017): Tensor Core 首次亮相,支援 FP16 輸入和 FP32 累加的矩陣乘加運算。在一個 SM 單元內,Tensor Core 可以每個時鐘執行 64 個 FP16 FMA(融合乘加)操作,使 V100 的深度學習峰值算力達到 125 TFLOPS(FP16),是同期 FP32 算力的 8 倍。Volta 的這一硬體創新直接催生了混合精度訓練的生產級方案。使用者必須在程式碼中顯式呼叫半精度資料型別,並手動管理損失縮放。
Turing 架構(T4/RTX 20 系列): 引入了 INT8、INT4 等整數格式的加速,但對 FP16 混合精度的增強相對有限。該架構更多面向推論場景的精度最佳化。
Ampere 架構(A100,2020): 這是混合精度生態的一次飛躍。A100 的第三代 Tensor Core 首次加入 BF16 和 TF32 的支援。TF32 是 19 位格式,在保持與 FP32 相同範圍的同時提供了 10 位尾數,方便使用者用最少的程式碼改動(僅需替換矩陣乘法函式)獲得 1.6 倍的加速。而 BF16 的硬體支援使得大型 Transformer 訓練可以告別繁複的動態損失縮放。此外,A100 的 Tensor Core 支援稀疏結構化矩陣乘,進一步將有效算力提升 2 倍。在 MLPerf Training 等基準測試中,A100 的混合精度訓練吞吐可達 V100 的 2~4 倍。
Hopper 架構(H100,2022-2023): 第四代 Tensor Core 帶來了 FP8 的原生支援,以及 Transformer Engine。Transformer Engine 是一個軟體-硬體協同最佳化庫,能夠動態地在 FP8 和 FP16/BF16 之間切換,並自動管理縮放因子。通過在矩陣乘法中使用 FP8,H100 的理論 FP8 Tensor Core 算力達到 3958 TFLOPS,是 A100 FP16 算力的 6 倍。更關鍵的是,Hopper 引入了塊級縮放機制:對於每個 128 元素的矩陣乘法輸入塊,可以附帶獨立的縮放因子,由硬體在乘加過程中自動解算。這解決了 FP8 直接表示梯度時極易產生的區域性上溢/下溢問題。在 Llama 3 70B 模型的訓練中,使用 H100 的 FP8 混合精度已實現約 1.4× 的額外吞吐提升,且模型質量無損。
除了 NVIDIA 路線,Google 的 TPU 系晶片從 TPU v2 開始就圍繞 BF16 建置生態,其 MXU(矩陣乘法單元)天然支援 BF16 輸入和 FP32 累加,使得基於 TPU 的訓練天然具有混合精度特性。華為 Ascend 系列 NPU 則支援 FP16 和 BF16,並通過達芬奇架構實現高效的半精度計算。AMD 的 Instinct MI300 系列也開始深度支援 BF16 和 FP8,追趕混合精度加速的潮流。
硬體的快速迭代不僅提升了混合精度的效率峰值,更降低了使用門檻。從 Volta 時代需要手動管理精度轉換和損失縮放,到 Hopper 時代通過 Transformer Engine 實現自動化精度決策,混合精度訓練正從一門“手藝”走向普適的基礎設施。
6. 前向傳播與啟用儲存的半精度化:頻寬與算力的雙重釋放
在混合精度訓練範式中,前向傳播是低精度算術紅利最直接的體現階段。該階段將計算密集的矩陣乘法替換為半精度 Tensor Core 操作,同時將啟用值的儲存精度降級,從而在算力和頻寬兩個維度釋放巨大潛力。
現代 Transformer 模型的前向傳播主要由線性層(投影矩陣乘法)和注意力機制中的矩陣乘法構成。以 Llama 3 70B 為例,一次前向傳播涉及的浮點運算次數約為 7×10¹² 次。若全部使用 FP32,即使 H100 的 FP32 向量算力也僅約 67 TFLOPS,完成一次前向需要超過 100 毫秒。但轉換為 FP16 或 BF16 Tensor Core 運算後,有效算力躍升至 990 TFLOPS 以上,前向時間可縮減至十幾毫秒。這種巨大的加速源於半精度矩陣乘法更短的指令執行週期和更高的並行度。
頻寬節省是另一項關鍵收益。大型模型訓練常常受限於視訊記憶體頻寬而非純粹算力。啟用值(每層 Transformer 的輸出)在訓練時必須保留用於反向傳播計算梯度,這佔據了視訊記憶體的大頭。在純 FP32 訓練中,啟用值儲存可能佔據總視訊記憶體使用量的 40%80%。將其轉換為 FP16 或 BF16,啟用值資料量立即減半。這不僅直接降低了單張 GPU 的視訊記憶體壓力,允許在相同硬體上使用更大的 micro batch size 或更長的序列長度,而且減少了 GPU 計算單元在等待資料時的空閒週期,提高了 MFU。以 Megatron-LM 訓練 GPT-3 175B 為例,通過啟用檢查點(Activation Checkpointing)和 FP16 啟用儲存的組合最佳化,單卡可容納的 micro batch size 提升了 24 倍,整體訓練吞吐提升了約 30%。
然而,啟用儲存的半精度化並非全無代價。某些啟用函式對於精度損失相對敏感。例如,Softmax 操作涉及指數和求和,FP16 的動態範圍可能導致指數上溢或求和精度不足,引起輸出分佈偏移。因此,典型的混合精度實現會在 Softmax 計算前將輸入臨時轉換為 FP32,計算完成後再截斷為 FP16 儲存。LayerNorm 和 RMSNorm 同樣需要 FP32 的累加器來保證小方差的準確計算。這些“高精度孤島”雖然略微增加了計算量,但保障了模型訓練的收斂特性。實踐中,架構會自動處理這些精度轉換,使用者通常無需干預。
此外,前向傳播還涉及 dropout、隨機深度等正則化操作。這些操作的資料流較小,通常浮點精度影響微乎其微,直接在半精度下進行即可。對於位置編碼(如 RoPE)這類涉及三角函式旋轉的操作,FP32 精度通常會被保留,以確保旋轉角度的精確性。
總體而言,前向傳播的半精度化是混合精度訓練中最成熟、收益最明確的部分。通過替換矩陣乘法運算元和智慧的精度管理,現代訓練系統可在不犧牲數值穩定性的前提下,持續獲取接近硬體峰值的計算效率。
7. 反向傳播中的梯度精度與頻寬節省:下溢之困與縮放之道
如果說前向傳播的半精度化主要釋放算力,那麼反向傳播的精度降低則更多關注頻寬節省和算力提升,但同時也引入了混合精度訓練中最棘手的問題——梯度下溢。
在反向傳播過程中,針對每一層的輸入梯度和權重梯度均通過鏈式法則計算得出。權重梯度的計算通常涉及啟用值的轉置與輸出梯度的矩陣乘法,這部分同樣是計算密集且頻寬密集的操作,可以從 Tensor Core 的半精度加速中獲益良多。輸入梯度則用於向前一層傳遞,同樣以半精度儲存和計算。整體上,反向傳播的浮點操作量約是前向傳播的兩倍,因此半精度加速帶來的絕對收益更加顯著。
然而,梯度的數值分佈特性與權重、啟用不同。隨著訓練進行,梯度往往呈現長尾分佈:大量梯度的絕對值極小,尤其在網路深層和訓練後期,部分梯度的量級可能低至 1e-10 到 1e-15 範圍,而 FP16 的最小正規數約為 6e-8,這意味著在沒有任何保護的情況下,大量有效梯度會被截斷為 0,導致引數無法更新。這種“下溢”現象在 FP16 中尤為嚴重,會表現為損失下降停滯或震盪,模型無法收斂到最佳狀態。
BF16 憑藉與 FP32 相同的 8 位指數,其最小正規數約為 1.17e-38,動態範圍覆蓋了絕大多數訓練場景下的梯度分佈,因此 BF16 訓練通常無須損失縮放,僅需在主權重更新時保持 FP32 即可。這也是 BF16 在大型模型社群迅速流行的核心原因之一。
損失縮放(Loss Scaling)是為 FP16 量身定製的解決方案。其原理簡單而有效:在計算損失之後、反向傳播之前,將損失值乘以一個較大的縮放因子 S(例如 1024)。由於反向傳播是線性過程,這一縮放會等比例地放大所有權重梯度和輸入梯度。當縮放後的梯度量級普遍進入 FP16 的可靠表示區間後,下溢問題便迎刃而解。在最佳化器更新之前,再將所有梯度除以 S,恢復到真實的梯度尺度。這樣,FP32 權重更新所接收到的梯度與未縮放時在數值上等價,但半精度梯度的低精度表示風險被規避。
一個典型的動態損失縮放策略會從初始 S 值(如 2¹⁶=65536)開始,每隔若干步檢查是否有梯度上溢(即出現 Inf 或 NaN)。出現上溢則跳過本次更新並減半 S;若連續數百步未出現上溢,則嘗試倍增 S 以更好地保護小梯度。這種自適應機制能在訓練的全程保持縮放因子的合理水平,成為 FP16 混合精度訓練中不可或缺的元件。
對於 FP8 訓練,梯度表示面臨更為嚴峻的挑戰。E5M2 格式的最小正規數約為 1.5e-7,仍然需要縮放;而 E4M3 的範圍更窄。因此 FP8 混合精度普遍採用逐張量或逐塊的縮放方法,對不同張量甚至不同 128 元素塊施加獨立的縮放因子,以保證數值穩定。這一技術細節正是 Hopper 架構 FP8 訓練高效且穩定的基石。
總體而言,反向傳播的精度選擇是混合精度訓練中精細權衡的集中體現:既要充分享受低精度帶來的頻寬和算力增益,又必須藉助損失縮放等技巧維護梯度的健康分佈。這一平衡一旦掌握,訓練的穩定性和效率便可兼得。
8. 損失縮放:理論與自適應演算法
損失縮放是 FP16 混合精度訓練的經典配套技術,其背後的數學原理簡潔而強大。理解其工作機理,對於調優混合精度訓練的穩定性和效率至關重要。
設原始損失值為 L,反向傳播計算得到的原始梯度為 ∇L。將 L 乘以縮放因子 S 得到縮放後的損失 L’ = S·L。由於微分運算元是線性的,縮放的損失產生的梯度 ∇L’ = S·∇L。因此,在反向傳播中,所有引數梯度都被放大了 S 倍。如果 S 選擇得當,這種放大將原始梯度從 FP16 的下溢區域移動到了安全的表示區間。最佳化器更新引數前,將梯度除以 S,得到還原的原始梯度值:∇L = (∇L’) / S。在 FP32 中進行這一除法操作是完全精確的,因此最終的引數更新與未縮放的情況在數學上等價。
這種“縮放-還原”的簡單技巧,之所以需要精心設計的自適應演算法,是因為 S 的選取面臨矛盾:S 越大,對小梯度的保護越好,但大梯度在上溢(超過 FP16 最大值 65504)的風險也越高。一旦出現上溢,梯度變為 Inf 或 NaN,會導致引數更新失敗,模型可能損壞。因此,必須動態地尋找 S 的上限,即剛好不會引發上溢的最大縮放因子。
動態損失縮放演算法通常遵循如下邏輯:設定一個初始縮放因子 S(如 2²⁴=1.67×10⁷),以及一個增長間隔和縮減因子。訓練過程中,每經過 N 次迭代(例如 2000 步),檢查期間是否有任何梯度的 FP16 表示中出現 Inf 或 NaN。若未出現上溢,則嘗試將 S 乘以一個增長因子(通常為 2),以更好地覆蓋更小梯度的新區域;若檢測到上溢,則捨棄本次迭代的梯度更新(不應用到模型),並將 S 除以縮減因子(通常為 2),並可能重置增長間隔。這種策略可以在不增加過多同步開銷的情況下,讓 S 自動適應當前訓練階段的梯度分佈。
一些改進的自適應演算法(如基於執行統計的損失縮放)通過記錄梯度大小分位數來預估合適的 S,使其更平滑地調整。此外,對於多 GPU 分散式訓練,需要在所有 GPU 之間同步上溢狀態,通常採用邏輯 OR 操作,即任一 GPU 出現上溢就降低全域性縮放因子。
與 FP16 相比,BF16 混合精度由於指數寬度與 FP32 一致,其典型的梯度下溢機率極低,因此在實踐中常採用固定縮放因子(例如 1.0,即無縮放)或簡單的靜態縮放因子(如 128),而無需動態調整。這簡化了訓練流程,也是 BF16 受到大型模型社群青睞的另一個重要原因。
在 FP8 混合精度時代,損失縮放的理念延伸為“張量級”或“塊級”縮放。由於 FP8 的範圍和精度都極其有限,單一的全域性損失縮放已不足以保證所有層的梯度健康。此時,每個矩陣乘法的輸入張量(或張量的一個子塊)都配有獨立的縮放因子,通過在執行中統計該張量的絕對最大值來確定合適的縮放幅度,確保資料對映到 FP8 的動態範圍中。Hopper 的 Transformer Engine 就內建了這一能力,軟體棧會根據張量統計自動插入縮放和反縮放操作,對使用者透明。
損失縮放的理論已非常成熟,但在實踐中,仍有一些微妙細節:縮放因子最好選擇 2 的冪次,以保證縮放和還原操作精確無損(只改變指數,不丟失有效位數)。動態縮放演算法的超引數(增長間隔、縮減因子等)需要針對模型和硬體微調,但好在預設值通常能在大多數場景下工作良好,這也是混合精度訓練能廣泛普及的一個前提。
9. 主權重副本與 FP32 引數更新:精度錨點的動力學
在混合精度訓練的整個迴圈中,FP32 主權重副本的存在是保障訓練最終精度與 FP32 基線一致的定海神針。這一設計解決了半精度尾數不足導致的引數更新“遺忘”問題,其動力學機制值得深入剖析。
現代最佳化器(如 Adam、AdamW)的引數更新量 ΔW 通常非常小。假設一個引數當前值為 0.1,梯度為 1e-5,學習率為 1e-4,則引數更新量 ΔW ≈ -1e-9。在 FP16 的表示中,0.1 的近似二進位制表示為 0.10009765625,下一個可表示的數約為 0.10015869140625,兩者間隔(即 FP16 在此區間的精度)約為 6.1e-5。更新量 -1e-9 遠小於這個間隔,如果直接加到 FP16 的權重上,結果將會被舍入回原始值,導致該次更新在數值上完全丟失。當學習率較小或梯度稀疏時,這種情況頻繁發生,模型將出現“更新停滯”現象。
FP32 主權重副本的運作方式破解了這一難題。任何一次半精度前向和反向傳播完成後,得到的半精度梯度先被還原(除以損失縮放因子)並轉換為 FP32。然後,這個 FP32 梯度被送入最佳化器,最佳化器基於 FP32 主權重和其內部的高精度狀態(Adam 的一階矩、二階矩通常也儲存在 FP32 中)計算出精確的 FP32 更新量。最後,這個更新量直接作用在 FP32 主權重上。由於 FP32 的精度高達約 7 個十進位制有效數字,對於 0.1 附近的更新步長可以精確到 1e-8 以下,確保了即使極微小的更新也能被累積。
比較形象的解釋是,FP32 主權重相當於一個“積分器”,可以忠實地累積來自半精度世界的每一次微小貢獻,而不會因為精度截斷丟失資訊。只有當需要下一輪前向傳播時,才從 FP32 主權重截斷出一份半精度副本供計算使用。這一截斷雖然會引入微小的表示誤差(FP32 → FP16 的四捨五入),但只要訓練損失的 landscape 不是過於病態,這種誤差通常不會影響收斂方向,因為它可以視為一種輕微的隨機擾動,反而可能幫助模型跳過某些尖銳的極小值。
FP32 最佳化器狀態的儲存是混合精度訓練視訊記憶體佔用中仍需 FP32 的部分。Adam 類的最佳化器需要儲存每個引數的一階矩 m 和二階矩 v,兩者通常以 FP32 格式儲存,這意味著最佳化器狀態佔用的視訊記憶體是模型引數本身的 8 倍(m 和 v 各 4 位元組)。這是混合精度無法進一步壓縮的部分,也因此催生了 8-bit 最佳化器(如 bitsandbytes 的 Adam8bit)的研究,它們在保持 FP32 主權重的同時,將最佳化器狀態量化為 8 位,進一步壓縮視訊記憶體。
此外,FP32 主權重副本的存在還在數值上提供了一種“保護帶”。在 FP16 混合精度下,由於尾數短,權重更新時的舍入誤差可能累積。FP32 主權重會定期(每個迭代)向下傳遞自己的精度狀態,重置傳播中的偏差。研究表明,若無 FP32 主權重,僅使用 FP16 進行完整的訓練,即使加損失縮放,模型精度往往會顯著劣於 FP32 基線(在 ImageNet 上 Top-1 準確率下降可能超過 1%)。而引入 FP32 主權重後,效能可以與基線持平甚至略有提升(可能由於半精度引入的適度噪聲有助於正則化)。
綜上所述,主權重副本與高精度引數更新是混合精度訓練將數值可靠性從“機率性”提升至“確定性”的關鍵發明,是連線低精度算術效率與高精度訓練完整性的橋樑。
10. 動態與靜態損失縮放策略:自動化精度管理的實踐
損失縮放的實施策略是混合精度訓練從實驗室走向工業大規模應用的重要工程細節。根據模型架構、數值格式和硬體特性,產業界形成了靜態縮放、動態縮放和自適應縮放三種主要路徑。
靜態損失縮放是最簡單的形式,即在整個訓練過程中使用恆定的縮放因子,或按預設的步數階梯式調整。對於 BF16 訓練,由於其動態範圍足夠大,可以安全地使用靜態縮放因子 1.0(等效於不縮放),或設定為 128、1024 等較小常數以防萬一。對於 FP16 訓練,若模型規模較小、訓練超引數穩定,也可以經過幾輪試跑後確定一個“安全”的靜態 S 值並固定使用。例如,某些 ResNet 類視覺模型使用靜態縮放因子 128 即可穩定收斂。靜態縮放的優點是實現簡單,沒有額外的分支判斷和同步開銷;缺點是缺乏靈活性,當訓練進入不同階段或更換資料集時可能需要手動調整。
動態損失縮放是當前 FP16 混合精度訓練的事實標準。其核心是一個自動機,根據梯度上溢的出現與否自動增加或降低 S。一個典型的動態縮放邏輯如上節所述。在 PyTorch 的 torch.cuda.amp.GradScaler 中,預設行為是初始縮放因子 2¹⁶(65536),每 2000 次迭代未出現 inf/NaN 則倍增 S,出現 inf/NaN 則跳過本次更新並將 S 減半。使用者還可以調整增長週期、回退因子和增長因子等。動態縮放的健壯性使其成為各類模型的預設選項,但它引入了兩個潛在成本:一是檢查 inf/NaN 需要遍歷梯度張量,略增計算開銷(通常可忽略);二是在分散式訓練中,所有 GPU 需要同步上溢標誌,增加了通訊步驟。不過,現代架構通常通過融合核函式將這一開銷降至近乎零。
自適應損失縮放是更高階的變種,它不再依賴簡單的“上升/下降”規則,而是維護執行時梯度量級的直方圖或統計資訊,基於分位數預估最佳縮放因子。例如,計算過去若干步梯度的 99.9% 分位數值,然後設定縮放因子使得該分位數對映到 FP16 最大值的 1/2 以內,這樣既留有餘量防止上溢,又儘可能大地保護小梯度。NVIDIA 的 APEX 庫中曾提供類似功能。這類方法更適合梯度分佈變化劇烈的訓練場景(如 GAN 訓練),能取得更穩定的縮放因子,但恆定計算和儲存開銷稍高。
在 FP8 訓練場景中,縮放策略從全域性走向細粒度。Hopper 架構的 FP8 混合精度訓練中,Transformer Engine 會動態計算每個 GEMM 操作輸入張量的最佳縮放因子。具體做法是:對於權重張量,由於其在整個訓練過程中不變(指前向計算時),縮放因子可以預先計算並儲存;對於啟用和梯度張量,則需在每個迭代中基於實際資料計算塊的絕對最大值,並應用縮放,這稱為“即時縮放”。這種細粒度縮放能將 FP8 有限的表示區間充分用滿,極大提升效率。
選擇合適的縮放策略需要結合硬體、模型和精度格式綜合考慮。幸運的是,主流架構已將這些策略高度自動化,研究人員通常只需選擇精度模式(如 torch.cuda.amp 或 torch.autocast),即可享受到先進的縮放管理,極大地降低了混合精度訓練的使用壁壘。
11. 效能評估核心指標體系:從吞吐到 MFU 的全景測量
評估混合精度訓練的實際收益,不能僅憑加速比的宣傳數字,而應建立一套多維度的指標體系,涵蓋吞吐量、視訊記憶體佔用、模型最終精度和硬體效率等關鍵維度。以下五個核心指標構成了衡量混合精度實施質量的事實標準。
1. 吞吐量加速比(Speedup)
定義為混合精度訓練的每秒處理樣本數(或每秒消耗 Token 數)與純 FP32 基線的比值。理想情況下,FP16/BF16 混合精度由於 Tensor Core 的算力優勢和頻寬節省,理論可達 2× 或更高加速。但在分散式大型模型訓練中,由於通訊瓶頸、資料載入和計算圖中的 FP32 高精度操作佔比,實際加速比往往略低。根據 MLPerf Training 和各大實驗室的公開技術報告,在數百至數千張 GPU 規模的 LLM 訓練中,FP16/BF16 混合精度相較於最佳化充分的 FP32 基線,吞吐加速比通常在 1.6× 至 2.15× 之間。FP8 混合精度則可在此基礎上再提升約 30%~40%。
2. 視訊記憶體佔用比
通過將啟用值和梯度的儲存精度從 FP32 降至 FP16/BF16,這部分視訊記憶體佔用理論減半。在實際大型模型訓練中,啟用值通常佔據總視訊記憶體的 40%~80%,因此混合精度可使總視訊記憶體需求降至純 FP32 的 50%~70%。啟用 FP8 後,對應線性層的啟用和梯度可進一步降至 25%~37.5%。視訊記憶體的節省直接轉化為支援更大的 batch size 或更長的序列長度,這對模型質量和訓練效率均有顯著正面影響。例如,訓練 Llama 3 時,使用 FP8 混合精度使單卡能處理的序列長度翻倍,某些團隊因此避免了更耗時的序列並行方案。
3. 模型最終精度差異
這是混合精度訓練能否被採用的“一票否決”指標。通常以驗證集困惑度(PPL)或下游任務準確率衡量。大規模的權威研究(如 Baziotis 等人 2024 年針對大語言模型 FP16/BF16/FP32 的對比分析)表明,在 NLP 大型模型上,BF16 基線與 FP32 基線的最終指標波動在 ±0.1% 以內,無統計學顯著差異。計算機視覺模型也有類似結論,甚至一些工作發現 FP16 訓練引入的適度噪聲對小模型泛化有輕微正面作用。只要訓練過程中動態損失縮放或 BF16 使用得當,最終精度損失是可控且幾乎可忽略的。這一結論得到了產業界大規模訓練的反覆驗證。
4. 動態損失縮放成功率
這是衡量 FP16 訓練穩定性的直接指標。它定義為訓練過程中未因梯度上溢而丟棄更新的步數佔總步數的百分比。業界共識是,若該比例 低於 90%,則說明動態縮放因子頻繁回退,訓練存在數值不穩定風險,可能暗中損害收斂。此時應檢查縮放初始值、調整策略或考慮遷移到 BF16。在健康的 FP16 訓練中,該比例通常達到 95% 以上,甚至長期維持在 100%。
5. 硬體峰值半精度利用率(MFU)
模型 FLOPs 利用率(Model FLOPs Utilization)是衡量硬體算力實際發揮程度的關鍵指標。計算公式為:實際吞吐下的有效 FLOPs / 硬體理論峰值半精度 FLOPs。在純 FP32 訓練中,由於頻寬限制和核心效率問題,MFU 常低於 30%。引入混合精度後,隨著 Tensor Core 的深度利用和頻寬需求降低,A100 上的 LLM 訓練 MFU 可提升至 50% 左右,而 H100 結合 FP8 和 FlashAttention-3 等最佳化後,MFU 已可突破 60%~70%(來源:Meta 訓練 Llama 3 405B 的工程報告,2024)。MFU 是綜合性指標,混合精度帶來的提升直接反映了端到端系統效率的最佳化程度。
這些指標應作為一個整體進行追蹤。單一的加速比高可能掩蓋了視訊記憶體上溢或模型精度下降,而單純追求高 MFU 可能忽略最終精度。優秀的混合精度實施方案應當在這五個維度上均達到產業基準水平,實現效率與精度的最優平衡。
12. 實際產業應用與架構支援:從研究原型到標準化 API
混合精度訓練從 2017 年的學術論文演變為當前深度學習架構的標準內建功能,僅用了不到三年時間。這一過程體現了產業對高效大型模型訓練的迫切需求與架構生態快速成熟的合力。
PyTorch 在 1.6 版本中引入了 torch.cuda.amp 模組,提供自動混合精度(Automatic Mixed Precision, AMP)支援。使用者只需在模型前向傳播程式碼外圍包裹 torch.cuda.amp.autocast 上下文管理器,並配合 GradScaler 即可完成 FP16 混合精度訓練的整合。AMP 會自動識別哪些操作應使用 FP16(如矩陣乘法、卷積),哪些應保持 FP32(如 Softmax、Normalization、Loss),從而在最小程式碼改動下實現高效混合精度。在 PyTorch 2.0 及以後,torch.compile 和 torch.autocast 進一步提升了自動化程度和效能。針對 BF16,PyTorch 也提供了相應的 dtype 支援,與 AMP 結合使用。
TensorFlow 從 2.4 版本開始,通過 tf.keras.mixed_precision API 提供混合精度。其設計理念與 PyTorch 類似,通過全域性策略(Policy)設定計算精度和變數精度,並自動插入損失縮放。TensorFlow 的 XLA 編譯器能夠進一步最佳化混合精度計算圖,在 TPU 和 GPU 上均能獲得良好加速。
NVIDIA APEX 是混合精度訓練的早期先鋒,提供了更細粒度的控制,如 O1(自動混合精度,部分操作保持 FP32)、O2(除了 BN 和損失外幾乎全 FP16)、O3(完全 FP16)等最佳化級別。AMP 模組的思想很大程度上借鑑了 APEX 的實踐。如今 APEX 中的部分高階特性(如多 GPU 通訊最佳化)仍在一些大型模型訓練架構中使用。
輝達 Megatron-LM 和 微軟 DeepSpeed 等大語言模型訓練架構深度集成了混合精度訓練。DeepSpeed 的 ZeRO 最佳化器將模型狀態分割到多卡,結合 BF16 混合精度訓練,實現了視訊記憶體和通訊的雙重最佳化,是訓練數千億引數模型的主流選擇。Megatron-LM 則針對 Transformer 結構對 FP16 和 BF16 做了精細的手工最佳化,包括自定義的損失縮放和精度轉換核函式。這些架構的混合精度支援經過大規模實測檢驗,可靠性極高。
JAX 和 Flax 生態則從底層提供了強大的自動向量化和精度管理能力。結合 jax.numpy 和 jax.lax 的 precision 引數,使用者可以靈活指定計算精度。Google 的 TPU 訓練棧幾乎全面採用 BF16,損失縮放需求天然消失,訓練複雜度進一步降低。
在產業實踐中,大規模訓練通常遵循“BF16 + FP32 主權重”的配置,免去損失縮放的調參負擔。例如,Meta 訓練 Llama 2 70B 時就採用了這一組合;而 NVIDIA 在使用 H100 訓練 Llama 3 時,則採用了 FP8 + BF16 的多層級混合精度,獲取極致的硬體利用率。架構側正在朝向完全自動化的方向演進:使用者只需指定目標精度或硬體型號,架構自動選擇最優的張量資料格式、縮放策略