如何解決混合精度訓練大模型的限制問題

混合精度已成為訓練大型深度學習模型的必要條件，但也帶來了許多挑戰。將模型參數和梯度轉換為較低精度資料類型（如FP16）可以加快訓練速度，但也會帶來數值穩定性的問題。使用進行FP16 訓練梯度更容易溢出或不足，導致優化器計算不精確，以及產生累加器超出資料類型範圍的等問題。

在這篇文章中，我們將討論混合精確訓練的數值穩定性問題。為了處理數值上的不穩定性，大型訓練工作經常被擱置數天，會導致專案的延期。所以我們可以引入Tensor Collection Hook來監控訓練期間的梯度條件，這樣可以更好地理解模型的內部狀態，更快辨識數值不穩定性。

在早期訓練階段了解模型的內部狀態可以判斷模型在後期訓練中是否容易出現不穩定是非常好的方法，如果能夠在訓練的頭幾個小時就能識別出梯度不穩定性，可以幫助我們提升很大的效率。所以本文提供了一系列值得關注的警告，以及數值不穩定性的補救措施。

混合精準度訓練

隨著深度學習繼續朝向更大的基礎模型發展。像GPT和T5這樣的大型語言模型現在主導NLP，在CV中對比模型(如CLIP)的泛化效果優於傳統的監督模型。特別是CLIP的學習文本嵌入意味著它可以執行超過過去CV模型能力的零樣本和少樣本推理，訓練這些模型都是一個挑戰。

這些大型的模型通常涉及深度transformers網絡，包括視覺和文本，並且包含數十億個參數。 GPT3有1750億個參數，CLIP則是在數百tb的圖像上進行訓練的。模型和資料的大小意味著模型需要在大型GPU叢集上進行數週甚至數月的訓練。為了加速訓練減少所需gpu的數量，模型通常以混合精度進行訓練。

混合精準訓練將一些訓練操作放在FP16中，而不是FP32。在FP16中進行的操作需要更少的內存，並且在現代gpu上可以比FP32的處理速度快8倍。儘管在FP16中訓練的大多數模型精度較低，但由於過度的參數化它們沒有顯示任何的性能下降。

隨著英偉達在Volta架構中引入Tensor Cores，低精準度浮點加速訓練更加快速。因為深度學習模型有很多參數，任何一個參數的確切值通常都不重要。透過以16位元而不是32位元來表示數字，可以一次性在Tensor Core暫存器中擬合更多參數，增加每個操作的並行性。

但FP16的訓練是存在挑戰性的。因為FP16不能表示絕對值大於65,504或小於5.96e-8的數字。深度學習框架例如如PyTorch帶有內建工具來處理FP16的限制(梯度縮放和自動混合精度)。但即使進行了這些安全檢查，由於參數或梯度超出可用範圍而導致大型訓練工作失敗的情況也很常見。深度學習的一些組件在FP32中發揮得很好，但是例如BN通常需要非常細粒度的調整，在FP16的限制下會導致數值不穩定，或者不能產生足夠的精度使模型正確收斂。這意味著模型並不能盲目地轉換為FP16。

所以深度學習框架使用自動混合精確度(AMP)，它透過一個預先定義的FP16訓練安全操作清單。 AMP只轉換模型中被認為安全的部分，同時將需要更高精度的操作保留在FP32中。另外在混合精度訓練中模型中透過給一些接近零梯度（低於FP16的最小範圍）的損失乘以一定數值來獲得更大的梯度，然後在應用優化器更新模型權重時將按比例向下調整來解決梯度過小的問題，這種方法稱為梯度縮放。

以下是PyTorch中一個典型的AMP訓練循環範例。

梯度縮放器scaler會將損失乘以一個可變的量。如果在梯度中觀察到nan，則將倍數降低一半，直到nan消失，然後在沒有出現nan的情況下，預設每2000步逐漸增加倍數。這樣會保持梯度在FP16範圍內，同時也防止梯度變為零。

訓練不穩定的案例

儘管框架都盡了最大的努力，但PyTorch和TensorFlow中內建的工具都無法阻止在FP16中出現的數值不穩定情況。

在HuggingFace的T5實作中，即使在訓練後模型變體也會產生INF值。在非常深的T5模型中，注意力值會在層上累積，最終達到FP16範圍之外，這會導致值無限大，例如在BN層中出現nan。他們是透過將INF值改為在FP16的最大值解決了這個問題，並且發現這對推斷的影響可以忽略不計。

另一個常見問題是ADAM優化器的限制。作為一個小更新，ADAM使用梯度的第一和第二矩的移動平均值來適應模型中每個參數的學習率。

這裡Beta1 和 Beta2 是每個時刻的移動平均參數，通常分別設定為 .9 和 .999。以 beta 參數除以步數的冪消除了更新中的初始偏差。在更新步驟中，為二階矩參數新增一個小的 epsilon 以避免被零除產生錯誤。 epsilon 的典型預設值是 1e-8。但 FP16 的最小值為 5.96e-8。這意味著如果二階矩太小，更新將除以零。所以在 PyTorch 中為了訓練不會發散，更新將跳過該步驟的變更。但問題仍然存在尤其是在 Beta2=.999 的情況下，任何小於 5.96e-8 的梯度都可能會在較長時間內停止參數的權重更新，優化器會進入不穩定狀態。

ADAM的優點是透過使用這兩個矩，可以調整每個參數的學習率。對於較慢的學習參數，可以加快學習速度，而對於快速學習參數，可以減慢學習速度。但如果對多個步驟的梯度計算為零，即使是很小的正值也會導致模型在學習率有時間向下調整之前發散。

另外PyTorch目前還一個問題，在使用混合精度時自動將epsilon更改為1e-7，這可以幫助防止梯度移回正值時發散。但是這樣做會帶來一個新的問題，當我們知道梯度在相同的範圍內時，增加ε會降低了優化器適應學習率的能力。所以盲目的增加epsilon也無法解決因為零梯度而導致訓練停滯的情況。

CLIP訓練中的梯度縮放

為了進一步證明訓練中可能出現的不穩定性，我們在CLIP影像模型上建立了一系列實驗。 CLIP是一種基於對比學習的模型，它透過視覺轉換器和描述這些圖像的文字嵌入同時學習圖像。對比組件試圖在每批資料中將影像匹配回原始描述。由於損失是在批次中計算的，在較大批次上的訓練已被證明能提供更好的結果。

CLIP同時訓練兩個transformers模型，一個類似GPT的語言模型和一個ViT圖像模型。兩種模型的深度都為梯度增長創造了超越FP16限制的機會。 OpenClip（arxiv 2212.07143）實作描述了使用FP16時的訓練不穩定性。

Tensor Collection Hook

為了更好地理解訓練期間的內部模型狀態，我們開發了一個Tensor Collection Hook (TCH)。 TCH可以包裝一個模型，並定期收集關於權重、梯度、損失、輸入、輸出和最佳化器狀態的摘要資訊。

例如，在這個實驗中，我們要找到並記錄訓練過程中的梯度條件。例如可能想每隔10步從每一層收集梯度範數、最小值、最大值、絕對值、平均值和標準差，並在 TensorBoard 中視覺化結果。

然後可以用out_dir作為--logdir輸入啟動TensorBoard。

實驗

為了重現CLIP中的訓練不穩定性，用於OpenCLIP訓練Laion 50億圖像資料集的一個子集。我們用TCH包裝模型，定期保存模型梯度、權重和優化器時刻的狀態，這樣就可以觀察到不穩定發生時模型內部發生了什麼事。

從vvi - h -14變體開始，OpenCLIP作者描述了在訓練期間存在穩定性問題。從預訓練的檢查點開始，將學習率提高到1-e4，與CLIP訓練後半段的學習率相似。在訓練進行到300步驟時，有意連續引入10個難度較高的訓練批次。

損失會隨著學習率的增加而增加，這是可預期的。當在第300步引入難度較大的情況時，損失會有一個小的，但不是很大的增加。該模型發現難度較大的情況，但沒有更新這些步驟中的大部分權重，因為nan出現在梯度中(在第二個圖中顯示為三角形)。通過這組難度較高的情況後，梯度降為零。

PyTorch梯度縮放

這裡發生了什麼事?為什麼梯度是零?問題就出在PyTorch的梯度縮放。梯度縮放是混合精度訓練中的重要工具。因為在具有數百萬或數十億個參數的模型中，任何一個參數的梯度都很小，並且通常低於FP16的最小範圍。

當混合精確訓練剛提出時，深度學習的科學家發現他們的模型在訓練早期通常會按照預期進行訓練，但最終會出現分歧。隨著訓練的進行梯度趨於變小，一些下溢的 FP16 變為零，使訓練變得不穩定。

為了解決梯度下溢，早期的技術只是簡單地將損失乘以一個固定的量，計算更大的梯度，然後將權重更新調整為相同的固定量(在混合精確訓練期間，權重仍儲存在FP32中)。但有時這個固定的量還是不夠。而較新的技術，如PyTorch的梯度縮放，從一個較大的乘數開始，通常是65536。但是由於這可能很高，導致大的梯度會溢出FP16值，所以梯度縮放器監視將溢出的nan梯度。如果觀察到nan，則在這一步驟跳過權重更新將乘數減半，然後繼續下一步。這一直持續到在梯度中沒有觀察到nan。如果在2000步驟中梯度縮放器沒有偵測到nan，它將嘗試將乘數加倍。

在上面的範例中，梯度縮放器完全按照預期工作。我們向它傳遞一組比預期損失更大的情況，這會產生更大的梯度導致溢出。但問題是現在的乘數很低，較小的梯度正在下降到零，梯度縮放器不監視零梯度只監視nan。

上面的例子最初看起來可能有些故意的成分，因為我們有意將困難的例子分組。但是經過數天的訓練，在大批量的情況下，產生nan的異常情況的機率肯定會增加。所以遇到足夠多的nan將梯度推至零的幾率是非常大。其實即使不引入困難的樣本，也常常會發現在幾千個訓練步驟後，梯度總是零。

產生梯度下溢的模型

為了進一步探索問題何時發生，何時不發生，將CLIP與通常在混合精度下訓練的較小CV模型YOLOV5進行了比較。在這兩種情況下的訓練過程中追蹤了每一層中零梯度的頻率。

在前9000步驟的訓練中，CLIP中5-20%的層顯示梯度下溢，而Yolo中的層僅顯示偶爾下溢。 CLIP中的下溢率也隨著時間的推移而增加，使得訓練不太穩定。

使用梯度縮放並不能解決這個問題，因為CLIP範圍內的梯度幅度遠大於YOLO範圍內的梯度幅度。在CLIP的情況下，當梯度縮放器將較大的梯度移到FP16的最大值附近時，最小的梯度仍然低於最小值。

如何解決解CLIP中的梯度不穩定性

在某些情況下，調整梯度縮放器的參數可以幫助防止下溢。在CLIP的情況下，可以嘗試修改以一個更大的乘數開始，並縮短增加間隔。

但是我們發現乘數會立即下降以防止溢出，並迫使小梯度回到零。

改進縮放比例的解決方案之一是使其更適應參數範圍。例如論文 Adaptive Loss Scaling for Mixed Precision Training 建議按層而不是整個模型執行損失縮放，這樣可以防止下溢。而我們的實驗顯示需要一種更具適應性的方法。由於 CLIP 層內的梯度仍然涵蓋整個 FP16 範圍，因此縮放需要適應每個單獨的參數以確保訓練穩定性。但是這種詳細的縮放需要大量記憶體會減少了訓練的批次大小。

較新的硬體提供了更有效的解決方案。例如BFloat16 (BF16) 是另一種 16 位元資料類型，它以精確度換取更大的範圍。 FP16 處理 5.96e-8 到 65,504，而BF16 可以處理 1.17e-38 到 3.39e38，與 FP32 的範圍相同。但是 BF16 的精度低於 FP16，會導致某些模型不收斂。但對於大型的transformers模型，BF16 並未顯示會降低收斂性。

我們執行相同的測試，插入一批困難的觀察結果，在BF16 中，當引入困難的情況時，梯度會出現尖峰，然後返回常規訓練，因為梯度縮放由於範圍增加而從未在梯度中觀察到NaN。

對比FP16和BF16的CLIP，我們發現BF16只有偶爾的梯度下溢。

在PyTorch 1.12及更高版本中，可以透過對AMP的一個小更改來啟動BF16。

如果需要更高的精確度，可以試試Tensorfloat32 (TF32)資料型別。 TF32由英偉達在安培GPU中引入，是一個19位浮點數，增加了BF16的額外範圍位，同時保留了FP16的精度。與FP16和BF16不同，它被設計成直接取代FP32，而不是在混合精度下啟用。要在PyTorch中啟用TF32，在訓練開始時加入兩行。

這裡要注意的是：在PyTorch 1.11之前，TF32在支援該資料類型的gpu上預設為啟用。從PyTorch 1.11開始，它必須手動啟用。 TF32的訓練速度比BF16和FP16慢，理論FLOPS只有FP16的一半，但仍比FP32的訓練速度快很多。

如果你用亞馬遜的AWS：BF16和TF32在P4d、P4de、G5、Trn1和DL1實例上是可用的。

在問題發生之前解決問題

上面的範例說明如何識別和修復FP16範圍內的限制。但這些問題往往在訓練後期才會出現。在訓練早期，模型會產生更高的損失並對異常值不太敏感，就像在OpenCLIP訓練中發生的那樣，在問題出現之前可能需要幾天的時間，這回浪費了昂貴的計算時間。

FP16和BF16都有優點和缺點。 FP16的限制會導致不穩定和失速訓練。但BF16提供的精度較低，收斂性也可能較差。所以我們肯定希望在訓練早期識別易受FP16不穩定性影響的模型，這樣我們就可以在不穩定性發生之前做出明智的決定。所以再次比較那些表現出和沒有表現出後續訓練不穩定性的模型，可以發現兩個趨勢。

在FP16中訓練的YOLO模型和在BF16中訓練的CLIP模型都顯示梯度下溢率一般小於1%，且隨著時間的推移是穩定的。

在FP16中訓練的CLIP模型在訓練的前1000步驟中下溢率​​為5-10%，並隨著時間的推移呈現上升趨勢。

所以透過使用TCH來追蹤梯度下溢率，能夠在訓練的前4-6小時內辨識出更高梯度不穩定性的趨勢。當觀察到這種趨勢時可以切換到BF16。

總結

混合精準訓練是訓練現有大型基礎模型的重要組成部分，但需要特別注意數值穩定性。了解模型的內部狀態對於診斷模型何時遇到混合精度資料類型的限制非常重要。透過使用一個TCH包裝模型，可以追蹤參數或梯度是否接近數值極限，並在不穩定發生之前執行訓練更改，從而可能減少不成功的訓練運行天數。

以上是如何解決混合精度訓練大模型的限制問題的詳細內容。更多資訊請關注PHP中文網其他相關文章！