AI 訓練加速原理解析與工程實務分享

這次分享將系統性的分析在 AI 模式訓練過程中的主要效能瓶頸，以及目前針對這些瓶頸的主要的加速方案和技術原理，並介紹百度智慧雲端在這方面的一些實踐成果。

今天的分享，主要包括三個部分：

首先介紹我們為什麼需要做 AI 訓練加速，也就是整體背景和出發點是什麼；

第二部分我們會系統性的分析實際訓練過程中的可能會遇到的效能瓶頸問題，然後針對這些問題，介紹目前主要的加速方案；

第三部分介紹百度百舸平台的 AI 訓練加速套件 AIAK-Training 在一些模型訓練加速上的實作效果。

1. 為什麼需要 AI 訓練加速？

在 AI 系統中，一個模型從生產到應用，一般包括離線訓練和推理部署兩大階段。

離線訓練階段，就是產生模型的過程，使用者需要根據自己的任務場景，準備好訓練模型所需的資料集，以及神經網路演算法。

演算法可以理解為是一個高度複雜的非凸數學函數，函數中包含許多變數以及參數。模型訓練的過程其實就是在學習神經網路模型中的參數。

模型訓練開始後，會讀取數據，然後送入模型進行前向計算，並計算與真實值的誤差。再執行反向計算得到參數梯度，最後更新參數。訓練會進行多輪的資料迭代。

訓練完成之後，我們會儲存訓練好的模型，然後將模型做上線部署，接受使用者的真實輸入，透過前向計算，完成推理。

因此，無論是訓練或推理，核心都是資料計算。為了加速運算效率，一般都是透過 GPU 等異構加速晶片來進行訓練和推理。

另外，從深度學習模型發展歷程來看，為了能夠持續突破模型的精確度上限，模型參數量其實在快速的膨脹。然而更大的參數量，就會帶來更大的計算複雜度。

下圖左側是摘自一篇公開的論文，從這篇總結裡，我們看到在 2010 年之前，模型的計算量大約 20 個月翻倍。在 2010~2015 年，常規模型計算每 5-6 個月翻倍。而在 2015 年之後，衍生了大模型訓練的趨勢，計算量增加 10~100 倍。

模型訓練對算力以及基礎設施的要求越來越高，訓練需要更多的算力，也需要更長的時間，這也導致了需要更多的資源成本。這裡我們列舉了一些論文或研究中公開的成本數據，反應了模型訓練的費用是非常高的。

因此，如何穩定的進行模型訓練，如何持續降本增效其實至關重要。

在這樣的大背景下，百度智慧雲端推出了百度百舸· AI 異構運算平台，目標是為 AI 場景提供軟硬整合的解決方案。透過 AI 計算、AI 儲存、AI 加速、AI 容器四層技術堆疊，滿足上層業務場景的需求。

• AI 運算層，提供了包含高效能的 GPU、以及崑崙等異構晶片資源，以及高效能的 RDMA 或 IB 網絡，以及自研的超級 AI 電腦 X-MAN 等；
• AI 儲存層，包括物件儲存 BOS 滿足資料湖儲存的需求、以及專為 AI 設計的高效能平行檔案系統 PFS；
• AI 加速層，包括資料湖儲存加速套件RapidFS，AI 訓練加速套件 AIAK-Training，AI 推理加速套件 AIAK-Inference；
• AI 容器層，也即是資源調度層，利用雲原生的技術能力，滿足 GPU、AI 作業等彈性調度層，利用雲原生的技術能力，滿足 GPU、AI 作業等彈性調度的需求。雲端原生 AI 的內容在我們上一期的技術公開課程有特別分享。

當我們考慮做效能加速的時候，第一個想到的可能是使用更好的硬體。

這會帶來一定的效能提升，但是大部分情況下可能並沒有充分發揮出硬體的運算能力，核心的原因就是訓練程式碼的執行效率並沒有調到最優或更優的狀態。

• 首先，框架為了相容於更多的需求，一般會提供一些較為通用的優化能力，更多的會關注易用性、生態建設等；
• 其次，演算法工程師在設計演算法時，核心的精力在如何提升模型的精確度，效能效率關注不足，某些情況下框架提供的一些通用最佳化能力，也都沒有利用。

因此，當我們決定使用某種模型演算法時，為了達到較好的資源效率和訓練效率，我們需要有意識的去最佳化。不過這裡也有很多的技術挑戰：

• 效能影響因素比較多，需要根據不同模型本身的運行特點進行分析，沒有完全固定的最佳化方案；
• 效能最佳化需要理解工程實現的原理，例如當我們去做異質晶片運算的最佳化，需要專業的異質研發經驗才能開展，技術門檻較高；
• 還有些情況，當我們做分散式訓練時，可能還需要結合叢集硬體和網路拓撲，來優化分散式訓練場景下的擴展性問題。

這些挑戰極大地影響了模型訓練性能的調優，因此我們推出了AIAK-Training 的加速套件，期望透過抽象易用性的介面降低優化成本，並透過軟硬協同的最佳化手段，來充分加速客戶在百度智慧雲端上的模型訓練效能。

2. 訓練表現開銷分析與加速方案

#在介紹AIAK-Training 具體效果之前，我們先介紹下訓練加速這個主題下關鍵的技術思路和方案原理是什麼樣的。

因為模型訓練優化本身就是一個軟硬體綜合的工作，技術堆疊相對比較複雜，今天的內容肯定沒辦法涵蓋全部細節，我們盡量把關鍵思路講到。

首先我們看目前的模型訓練方案。在過去的發展階段裡，模型訓練方案關鍵有兩個層次的變化，一是從單卡訓練到分散式訓練的變化，二是從資料並行訓練到多維混合並行訓練的變化。這裡的核心驅動點一個是訓練的資料量，一個是模型的參數量。

• 單卡訓練方式：實際採用此模式，一般都是模型參數量和資料量相對比較少，單卡的訓練時間可以接受。模型參數規模需要確保在訓練過程中，單張卡的記憶體能夠滿足儲存的上限。依照新的 GPU 卡的記憶體容量配置，一般可放置的最大規模在10 億級參數量；
• 當資料集規模比較大的時候，因為訓練過程中需要多次遍歷全量的資料集，單卡訓練就會需要比較長的時間，這時可以從單卡擴展到多卡，透過分散式的方式來加速訓練。

這裡應用最廣泛的就是資料並行。在資料並行方案下，資料集會被平均切分成多份，然後每張卡上保存完整的模型，各自獨立並行處理切分後的子資料集。

當模型參數量夠大的時候，例如參數量達到百億、千億級別，單卡放不下完整的模型，這裡又出現了模型並行、或者同時使用數據並行和模型並行的混合並行方案。

模型並行，會將模型切分到不同的卡片上，每個卡片上放置模型的一部分，這裡又根據切分的方式不同，例如層內切分或層間切分，又細分了Tensor 並行和管線並行的方式。

因為在一般模型訓練中，使用更多的還是資料並行，我們下面重點還是以資料並行為例，來介紹效能最佳化的想法。

我們從軟硬體整體視角先理解下單卡訓練過程中存在的效能開銷。

下圖左邊是我們從軟體角度上的訓練流程。單卡訓練的過程，主要包括資料讀取、資料預處理，前向計算輸出、計算 loss，根據 loss 函數反向計算，得到每一層參數的梯度，最後根據梯度更新模型參數。持續該過程，直到訓練收斂。

下圖右邊，是一個簡化的節點硬體拓樸圖。最上面是資料存儲，可以是本地存儲，也可以是網路存儲。接著是 CPU、內存，CPU 下透過多個 PCIe Switch 連接 8 張 GPU 卡，編號從 0~7，8 張卡之間透過 NVSwitch 互聯。不同的計算實例，硬體的拓樸結構會有不同。

• 當訓練啟動時，首先資料讀取，涉及從儲存媒體讀取資料到記憶體中，主要是儲存I/O 開銷。根據儲存介質，要讀取的資料量不同，這部分時間開銷也不同；
• 接下來是資料預處理，主要是對讀入的資料進行一些資料增強的操作，例如對圖片進行resize，對比、飽和調整等等，這部分工作大多數情況下是在 CPU 上。當在 CPU 上完成資料預處理作業之後，需要從主機記憶體拷貝到 GPU 顯存上，涉及主機和裝置之間的記憶體拷貝開銷；
• 然後開始前向計算、反向計算、參數更新，這部分更多操作主要是透過 GPU 來進行，主要時間花費在 GPU 計算上面。這個過程裡也可能會穿插一些 CPU 上的操作，可能也需要做主機和裝置記憶體之間的拷貝。

因此，從單卡角度來看，主要存在 I/O、CPU 預處理、CPU 和 GPU 之間資料拷貝，GPU 運算等方面的開銷。

接著我們看下資料並行的過程。

下圖左邊還是訓練的主流程，右邊展示了一個 3 機 24 卡的訓練集群的硬體拓撲，3 台機器通過網路互聯。

前面的部分我們也介紹到了，在資料並行裡每個設備並行獨立地執行前向和反向計算過程，因此，每個訓練進程也會遇到前面講的單卡訓練中的效能開銷問題。

資料並行為了保證和單卡訓練在數學上等價，需要確保每張卡的模型參數在迭代過程中始終保持一致。這裡一方面需要讓各 GPU 卡的模型參數初始化狀態一致，這個一般是在訓練開始前，透過廣播的方式將第一張卡片上的參數狀態廣播到其他的卡片。

而在訓練期間，由於每個設備處理的資料不同，前向計算所得到的模型損失值也是不同的，因此還需要在每個設備反向計算出梯度之後，進行梯度的平均，以平均後的梯度值來更新模型參數，從而確保每張卡的模型參數在迭代過程中始終保持一致。

梯度平均涉及到通訊的過程，包括節點內部卡之間的通信，以及跨節點的網路通訊開銷。這裡的通信又包括同步通信和非同步通信，不過為了確保模型訓練的收斂，一般都是採用同步通信的方案，後續的優化工作也都是基於同步通信的方式來展開。

由上可知，資料並行相比單卡訓練，主要增加了額外的通訊開銷。

透過前述分析，我們知道加速AI 訓練不單是某一方面的工作，需要從資料載入、模型計算、分散式通訊等系統維度綜合考量。這裡說的資料加載，包括資料 I/O、預處理、記憶體拷貝等過程。

在具體的最佳化實踐中，給定一個待優化的模型，加速模型訓練就是要不斷提升訓練的總體吞吐（每秒可以訓練的樣本數）。這個過程，我們一般可以先分析單卡的訓練吞吐，當單卡的訓練吞吐提升上來之後，我們再看從單卡擴展到多卡，看看如何提升多卡的訓練加速比。

首先單卡的訓練優化，極限優化目標是全部時間都在 GPU 運算上，加速器利用率100%。當然實際上很難完全達到這個狀態，但是我們可以以這個指標來牽引或衡量我們的工作。

單卡效能最佳化關鍵包括兩部分：

• 首先資料載入效率的最佳化。從儲存系統上，我們可以使用更高效能的儲存介質，或是基於這些高速儲存媒體所組成的平行檔案系統，或是說一些快取加速系統。前面介紹到的，百度百舸也提供了對應的儲存系統方案，例如 PFS、RapidFS 等。除此之外，還需要在框架 dataloader 中最佳化資料的讀取過程。
• 其次就是模型計算效率的最佳化。主要考慮如何優化計算實現，怎麼提升計算單元的利用效率，這塊可能需要結合模型具體分析。

然後從單卡擴展到多卡，目標是如何達到線性加速比。線性加速比這個指標，簡單來說就是從 1 張卡擴到 2 張卡訓練時，訓練的表現是否是單卡的2倍。

這裡核心就是優化分散式的通訊效率，一方面是硬體層面的最佳化，另一方面在實際通訊中，需要考慮怎麼利用好網路的頻寬資源，或是是否能夠將通訊流程隱藏等。

下面我們分別從這幾個面向詳細展開。

首先是資料載入方面的最佳化。

當我們實例化一個 dataloader 之後，我們會持續迭代 dataloader 來讀取一個 batch 的資料進行模型訓練。

如果這裡不做任何最佳化，如下圖上半部所示，每個 batch 的資料載入過程和每個 batch 的模型訓練過程，實際上是串列進行的。從 GPU 觀點來看，就會出現因為資料載入所導致的運算間隙，運算資源存在時間上的浪費。

除了前面講到我們用更好的硬體直接提升資料讀的效率之外，還可以怎麼優化呢？

實際上在AI 訓練過程中，資料存取上有兩個關鍵特徵：

• 當資料集做完 shuffle 之後，每輪訓練所需的 batch 資料以及存取順序是已知的；
• 任兩個batch 的資料讀取可以並行，因為資料之間沒有任何依賴關係。

因此，我們在不依賴硬體層面改動的時候，可以做的一個優化工作就是數據預取，當訓練第一個batch 數據的時候，可以提前加載下一個 batch 數據，讓 I/O 的過程和GPU 上的計算充分並行。

首先，我們需要利用好dataloader 中已有的最佳化方案，一是合理設定 num_workers 超參數，透過多進程的方式讀取數據，這一步可以實現數據從儲存系統中預取到主機記憶體。二是從主機記憶體拷貝到 GPU顯存，可以透過 pinned memory 的機制來加速。

大概介紹下 pinned memory 加速的主要原理：記憶體資料有兩個類型 pageable memory 和 pinned memory，pageable memory 中的資料有被換出到磁碟上的可能。在這種情況下，執行 H2D 時，可能需要先從磁碟讀到內存，然後從內存拷貝到顯存。另外，pageable memory 資料拷貝到 GPU 顯存時，需要先建立一個臨時的 pinned memory 緩衝區，把資料從 pageable memory 拷貝 pinned memory，之後才能傳送到 GPU 上，也多了額外的資料搬運的操作。

不過當我們啟用了上述方案後，我們僅實現了從儲存系統到主機記憶體的預取，加快了主機到裝置的資料拷貝速度。但主機到裝置的記憶體拷貝，且實際計算 kernel 在 GPU 上還是序列執行，也就是 GPU 上仍存在少量的時間間隙。

AIAK 針對這個問題，又做了進一步的最佳化，可以實現 H2D 和前向計算的 overlap。

在資料並行場景下，還需要注意的一個事情，資料需要均衡的切分。

如果每個訓練進程分配的資料不均衡，計算量就會不同，也就導致每個進程前向計算和反向計算完成的時間不同，那麼先完成計算的進程，在反向過程中就會先進入到梯度通訊環節中，但是因為Allreduce 通訊是同步通訊操作，需要所有進程同時開始並同時結束，因此先開始通訊的進程，會一直等待其他所有進程也發起了 AllReduce 後才能一起完成通訊操作。這裡就會出現因為快慢不一導致的資源空閒問題。

為了解決這個問題，需要每個行程使用相同的 batchsize 來讀取數據，另外每個 batch 的資料量要均衡。圖像類別資料一般會固定尺寸進行訓練，而像 NLP 類別模型需要處理變長的語句，可能需要進行特殊的處理，例如可以將資料 padding 到相同的長度，或透過樣本長度排序的方式來均衡分配等。

以下介紹計算效率的最佳化。

計算包含前向、反向、參數更新。最佳化運算的目標，是為了能夠充分發揮異質硬體的算力，理想情況就是讓 GPU 晶片實際運算時的效能達到理論高峰。

我們先從一個單一算子的角度分析，當我們準備在 GPU 上執行一個運算運算的時候，簡化的流程有四個步驟。

• 首先在 CPU 上非同步發射一個GPU 計算Kernel；
• 當Kernel 調度執行時，需要先從 GPU 上的 Global Memory 讀取計算所需的資料;
• 開始執行計算;
• 當計算完成後，需要將計算的結果寫回 Global Memory。

根據計算、訪存的開銷佔比，一般會將算子分類為計算瓶頸或訪存瓶頸。

當從一個算子擴展到一個完整的模型訓練時，因為要連續執行非常多的計算Kernel，那麼 Kernel 計算之間就會出現很多因為 Kernel Launch、中間結果的讀寫導致的計算間隙問題。

由上可知，最佳化模型計算效率，需從訪存最佳化、計算最佳化、其他開銷的最佳化綜合考量。

• 訪問最佳化，主要考慮如何減少資料在顯存和計算單元之間搬運的時間。從算子實現角度上，需要利用好 GPU 儲存層次架構，例如把資料搬運到更快的記憶體上例如 share memory，減少對 global memory 的訪問，從而節省存取時間。或做好計算指令和存取指令的 overlap，來提升計算訪存比。而從單算子擴展到多算子時，還需要考慮如何減少中間結果的讀寫，這種問題一般可以用算子融合的手段來優化；
• #計算優化，計算瓶頸問題多半應該是沒有正確對任務進行分塊，或是沒有利用好 GPU 並行計算的優勢，導致並行度不高。還有可能沒有使用合併指令集導致運算效率低下，或沒有使用 Tensor Core 等高效能運算單元，造成資源浪費。不同的問題，也對應著不同的最佳化手段；
• Kernel Launch 等其他開銷，大量時間花費在訪存或計算之外，可以考慮的最佳化手段如算子融合、Cuda Graph 等。

首先是算子融合。算子在底層 GPU 執行時，會發起一次或多次 Kernel Launch，Kernel 之間互動資料也需要經過顯存，而算子融合就是將多個 GPU Kernel 融合成一個大 Kernel，統一發起和執行。

• 因為減少了需要執行的算符數量，從而可以減少 Kernel 調度和發起的開銷；
• 透過融合，可以透過暫存器等來傳遞中間結果，避免從 global memory 的來回搬運，大幅降低了顯存等待的時間；
• 在某些場景中，可以透過運算子融合，可以更充分的利用運算資源，提升資源和運算的效率。

算符融合具體要如何實現？

一種方式，分析模型中的低效率操作，專家經驗手寫融合算符。在 GPU 上主要就是 CUDA 算子研發，這裡有一定的門檻。 AIAK-Training 會針對典型的模型結構，或是顧客需求，提供高效率最佳化的算子實作。

另一種方式，就是編譯最佳化的方案。透過編譯的方式進行計算最佳化，以及程式碼自動生成，從而降低在不同硬體上的手動最佳化成本。不過目前很多編譯方案比較多還是針對推理優化，訓練上的方案還在快速演進過程中。不過從極致性能角度來看，未來一段時間依然離不開手寫融合算子的工作。

以下介紹一些算子融合的實際案例。第一個是針對典型模型網路結構的最佳化。

下圖展示了我們在 SwinTransformer 模型中針對核心模組 WindowAttention 進行的運算融合最佳化。

WindowAttention 結構，核心操作公式如下圖所示。計算過程中，需依序執行 7 個計算 Kernel。再加上一些 reshape 等轉換操作，總共需要 launch 10 個 Kernel。透過效能分析發現，實際執行過程中 launch kernel 的間隔冗餘開銷占到了端到端 80% 以上的時間，導致該模組存在著較大的最佳化空間。

透過將這些 Kernel 融合成一個，整個模組的執行時間從 392 微秒減少到 13 微秒，單算子加速了 30 倍。整個模型的訓練效率，端到端加速了 20% 以上。

這個最佳化的核心思路，主要有三點：

• 是利用好 GPU 的三級訪存流水：顯存、share memory、暫存器；
• 透過分塊策略，將2 個矩陣乘法和 softmax 進行融合；
• 針對前向中間結果的最佳化，充分利用Tensor Core，在反向計算中使用重計算代替重加載，使得訪存開銷極大降低。

下圖是一個資料運算的融合舉例，是在 FCOS3D 模型中對於座標壓縮運算的一個最佳化。

透過效能分析發現，這個操作過程中存在大量 GPU 空隙，GPU 利用率較低。此操作的主要功能是根據 index 將 3D-Tensor 壓縮成 2D-Tensor。在原生實作中會先在host 端產生 index，再進行 H2D 拷貝，最後完成 Tensor 壓縮，這會造成額外的拷貝和等待開銷。

為此，我們重新實現這部分操作，核心思維主要就是將操作全部遷移至GPU 上，直接在GPU 上完成 index 的生成和Tensor 的壓縮，減少 CPU 參與，同時避免了非必要的 CPU-GPU 之間的記憶體拷貝。

單算子執行時間從 9.69 毫秒減少到 32 微秒，加速了 300 倍，整個模型端到端訓練提升了 10% 以上。

下面我們介紹另一個運算最佳化的思路，就是提高運算的平行度，充分利用 GPU 並行運算的優勢，也藉助一些實際案例來介紹。

我們發現在一些模型中，有一些操作是串列執行的。例如在一些目標偵測模型裡，在loss 計算過程中，有些操作並不是按照一個 batch 進行操作，而是 for-loop 每張圖片或一個樣本，這種情況下，當我們去提升 batchsize 的時候，因為這裡的串行，可能沒辦法達到我們想要的效能效果。

以 YOLOv7 中的 SimOTA 操作舉例，在原生實作中，是透過 for-loop 遍歷一個 batch 的每一張圖片，然後為圖片的 gtbox 執行 SimOTA 標籤分配。這種串列的實作方式導致該部分操作的 GPU 利用率非常低效。

而每張圖片處理時資料之間是沒有依賴的。因此，我們所做的一項工作就是將串列計算改成 batch 並行計算，透過將一個 batch 的資料並行化的標籤分配，加速這部分計算的效率。

最終效果上，SimOTA 操作的耗時從384 毫秒下降到 69 毫秒，計算效率提升 5.5 倍，整個模型的端到端訓練效率提升了 18%以上。

模型訓練過程中也有其他的類似場景，例如說參數更新。參數更新時，預設也是透過循環的方式，遍歷每個參數，然後每個參數都會啟動一個參數更新的 Cuda Kernel，然後依序執行。

針對這種情況，AIAK 也增加了 FusedOptimizer 的最佳化，透過融合參數更新的算子，實現批量更新參數，大幅減少 Kernel Launch 次數。

下面介紹了另一個最佳化手段 CUDA Graph，主要是為了減少 CPU Launch Kernel 的開銷。

CUDA Graph 是在CUDA 10 版本中引入的特性，可以將一系列 CUDA Kernel 封裝成單一單元，可以透過一次 CPU Launch 作業啟動多個 GPU Kernel，從而減少了CPU Launch Kernel 的開銷。

如下圖所示，預設情況下 CPU 需要依序發射多個 Kernel，如果 Kernel 運算時間比較短，Kernel 之間的 Launch 間隙可能會成為效能瓶頸。透過 CUDA Graph，只需花費一些額外時間建構 Graph，後續透過一次 Graph 的發射，即可大幅縮短實際執行時 Kernel 之間的間隙。

現在許多框架也增加了對 CUDA Graph 的支持，透過插入一些程式碼來讓K能這個功能。不過也有一些使用上的限制，例如不支援動態 shape、不支援動態控制流程、過程中不能捕捉 CPU 操作等等，可以根據模型情況，試著使用下這種最佳化能力。

以下介紹最後一個計算最佳化手段，充分利用 Tensor Core 計算單元。

一個 GPU 內一般包含多個 SM，每個 SM 包含不同資料型別的運算核心，以及各類儲存資源等。下圖左邊所示，是一個 NVIDIA A100 SM 的示意圖，一個 SM 包含 64 個 FP32 CUDA Core，以及 4 個 Tensor Core。

模型訓練時，預設主要是用到FP32 CUDA Core 計算，而Tensor Core 是一塊特殊的硬體執行單元，是從 Volta 系列 GPU 開始引入，主要是用來加速矩陣或卷積的操作效率。

比起 FP32 CUDA Core 一次只能執行兩個標量的計算，Tensor Core 可以一次對兩個矩陣執行計算，因此 Tensor Core 的計算吞吐遠高於 FP32 CUDA Core。

在 A100 中，Tensor Core 支援了多種浮點資料類型，對於深度學習訓練來說，可能涉及包括 FP16、BF16、以及 TF32 模式。

TF32 主要用於單精確度訓練場景，相較於 FP32 訓練，在維持相同的訪存頻寬需求下，理論計算吞吐量提高了 8 倍。

FP16/BF16 主要用於混合精準度訓練場景，相較於 FP32 訓練，訪存需求減少了一半，理論計算吞吐量提高了 16 倍。

使用 Tensor Core 的方式，可以使用底層的 cublas 或 cuda 介面進行編程，而對於演算法開發者來說，則更直接的是使用框架中提供的 TF32 訓練或混合精確度訓練方案。

首先是 TF32 訓練模式，TF32 則是 Ampere 開始引進。

TF32 在浮點數的表達中，有 8 個指數位，10 個尾數位和 1 個符號位。指數位與  FP32 相同，即資料表示範圍相同，但尾數位低於 FP32，和 FP16 相同。

要注意的是，TF32 不是對外開放的數值類型，只是 Tensor Core 的計算模式，也就是使用者不能去直接建立一個 TF32 類型的浮點數。

當啟用 TF32 的時候，Tensor Core 計算矩陣或捲積作業時，會自動將 FP32 轉換成 TF32，計算完成後，輸出的資料型別仍是 FP32 型別。

TF32 訓練在某些框架版本中是預設開啟，某些框架版本中可能需要透過環境變數或參數配置來手動開啟，具體需要參考框架的使用者手冊。

不過由於 TF32 相較於 FP32 來說，精確度範圍降低了，實際訓練時也需要注意對模式收斂精確度的影響。

混合精準度訓練是指在盡可能減少模型精準度損失的情況下，使用 FP32 和 FP16 混合精準度進行訓練。

混合精準度訓練的效益主要有：相較於 FP32 訓練，記憶體需求減少，可以訓練更大的網路或使用更大的 batchsize。使用更少的記憶體頻寬，可以加速資料傳輸，半精確度的計算也可以讓數學運算效率更快；

#不過因為FP16 的指數位和尾數位的範圍都比 FP32 少，因此數值表示範圍和精度都會有降低，在實際使用的時候，就可能出現因為表示範圍狹窄導致的數值溢位問題，或是因為精度不足導致舍入誤差。

為了優化類似問題，混合精準度訓練方案中有幾個關鍵的技術工作：

• 算子黑白名單機制，框架使用黑白名單自動為算子選擇精確度，模型訓練過程中，會自動的插入 cast 作業進行類型轉換，不需要開發者介入。針對數值精度敏感的計算，仍使用 FP32 來算，而數值安全的計算，例如矩陣乘，則會使用 FP16 計算；
• 訓練過程中，會儲存一份FP32 的權重參數，用於訓練時候的參數更新，優化舍入誤差的問題；
• 針對 FP16 容易溢出的問題，使用了Loss scaling 的方案，例如對 Loss 放大n 倍，根據鍊式法則，梯度也會隨之放大n 倍，使其落到 FP16 的表示範圍內，具體過程如下圖左側所示。

目前所有框架都已經支援混合精度。 AIAK-Training 元件進一步引入 NVIDIA Apex 中  AMP O2 混合精確度模式，這個模式會更激進的將更多計算轉 FP16 來加速訓練。相較於預設 O1 模式，速度會有進一步提升，不過精確度可能會受影響，需要結合具體模型驗證。

AIAK-Training 提供相容 torch amp 原生用法的使用方式，方便啟用 O2 模式。

下面介紹通訊優化，這也是一個非常大的主題，牽涉到內容也非常多。

前面也介紹了，通訊主要是分散式訓練中引入的，因為從單卡擴展到多卡，需要進行多卡之間的一些資料同步，這些同步操作就是透過通訊來實施的。

下圖列了一個通訊優化的整體架構：

• 最底層就是網路協定層，包括傳統的 TCP 網路、以及在訓練場景用的越來越多的RoCE 或 IB 的高效能 RDMA 網路。這些底層網路方案，百度百舸也都提供了支援。顯然透過改善硬體基礎設施來提升網路頻寬、降低延遲，是最直接有效的優化通訊效能的方法。
• 然後是通訊庫層。因為需要使用底層的網路協定來進行通信，並且實際應用時可能涉及多種通信原語，例如點對點通信、集合通信等，因此也出現一些高度封裝並且優化的通信庫。在 GPU 訓練場景中，我們通常使用 NCCL 通訊庫，效能比較優。
• 基於底層通訊庫，上層框架可以比較方便的建構分散式訓練的通訊架構。常見的包括參數伺服器架構，集合通訊架構。目前 CV/NLP 等領域的模型，主要都是採用集合通訊的架構方式。
• 通訊策略層，主要是在應用層上做一些通訊效率的最佳化，例如通訊隱藏、通訊融合、通訊壓縮、通訊降頻等不同的思路，這些方式大部分也可以疊加一起使用。

我們先來看通訊策略層面的最佳化思路，首先是通訊隱藏優化。

在資料並行中，梯度同步通訊是在訓練的反向過程中進行的，當反向算出梯度之後，就可以進行全域的梯度平均。

如果不做任何的機制最佳化，反向計算和通訊就會串行的進行，就會存在計算上的時間間隙。

由於反向過程中，上一個梯度的通訊和下一個梯度的計算，兩者之間沒有任何資料依賴，因此可以讓上一個梯度通訊和下一個梯度計算並行起來，讓兩者的耗時相互重疊，從而可以隱藏部分的通訊耗時。

在實現層面上，通常是將通訊和計算算子調度到不同的cuda 流上實現的，通訊算子調度到通訊流，計算算子調度到計算流，不同流上的算子可以並行發射執行，從而實現反向中梯度通訊和計算的平行重疊。

目前這個最佳化能力，在框架中都是預設開啟的。

其次，通訊融合最佳化。

預設情況下，模型中的每個梯度都需要發起一次通訊操作。如果單一梯度的 size 比較小，那麼小資料包在實際通訊時，網路頻寬的使用率就會非常低，通訊效能較差。

通信融合，就是將多個梯度融合到一起統一進行一次通信，從通信開銷模型上分析，既能提升頻寬利用率，又能減少通信的初始化延遲項。

現在很多分散式訓練框架中，也預設都支援梯度融合的策略，不同框架實現方式有一定的區別，有的實現需要先進行梯度協商確定通信順序，有的則是直接通過靜態的通信分桶。

雖然框架中預設支援通訊融合，但是梯度融合的大小一般可以透過參數來配置，使用者可以根據物理環境和模型的需求，調節合適的融合閾值，應該可以取得更佳的收益。

如果在網路頻寬較低的訓練場景中，例如低頻寬的 TCP 環境中，梯度同步的延遲可能會成為訓練的主要效能瓶頸。

• 這種情況下，可以考慮的一個最佳化手段就是通訊壓縮。通訊壓縮，主要有三種不同的壓縮想法：
• 量化壓縮，例如使用更低精度來表示梯度，這種方式的壓縮率較低，最大能從 32 位元壓縮到1 位，也就是最大壓縮32 倍。
• 稀疏化壓縮，典型的演算法例如 DGC 演算法，核心思想是每輪迭代只傳輸重要的梯度，也就是梯度數值超過設定的某一個閾值，同時為了減少資訊損失，會把剩下下不重要的梯度在本地累積，只要時間足夠，最終累積梯度就會超過所設定的閾值，再進行梯度交換。透過這種方式減少通訊的資料量，降低網路頻寬的需求。同時為了減少對模型收斂的影響，也透過動量校正、梯度裁切、動量因子遮蔽、預熱訓練等不同方式來緩解。這種方案目前主要支援 SGD 優化器。
• 低秩矩陣壓縮方式，典型的演算法例如 PowerSGD，核心想法是將一個大的梯度矩陣分解成多個小梯度矩陣，透過傳輸小矩陣來減少通訊量。

通訊降頻優化，最簡單的思路就是增大batchsize，每次迭代更多的數據，減少了迭代次數，也就是減少了通信量。

不過 batchsize 也不是越大越好，越大的 batchsize 可能會導致模型收斂精確度下降或收斂速度變慢。針對類似問題，業界也提出例如 LARS、LAMB 優化器演算法，透過分層自適應調整學習率，緩解類似問題，AIAK-Training 也增加了支援。

為了增加 batchsize，如果顯存比較充足，可以直接調整 batchsize 超參。如果顯存比較吃緊，還可以透過梯度累加的方式，跳過若干次梯度通信，實際上也相當於增加了 batchsize。

下面介紹了一個針對通訊拓撲的最佳化方案－分層拓樸通信，也是針對機間網路頻寬比較低的情況。

透過分層通信，可以充分的利用機內的高的互聯頻寬，同時弱化機間低網路頻寬的影響。

AIAK 中也實現了這個通訊方案，在 25Gbps TCP 環境下，實測 4 機 32 卡 SwinTransformer 訓練，透過分層 allreduce，效能可以加速 85%。

最後介紹一個底層通訊庫層面的最佳化，GPU Direct RDMA 通訊技術，這個技術需要硬體環境上支援 RDMA 網路。

RDMA 通信，允許本地應用程式直接讀寫遠端應用程式的用戶態虛擬內存，整個通信過程，除了最一開始提交發送請求這一步需要CPU 的參與，其餘都是由網卡硬體完成的，不需要記憶體拷貝、系統中斷以及軟體處理，因此可以實現極致的低時延和高頻寬。

而在 GPU 場景中，GPU Direct RDMA 技術進一步增加了 RDMA 直接存取 GPU 顯存的支持，避免通訊的時候，資料在 GPU 記憶體和主機記憶體中來回拷貝，跨機通訊延遲進一步降低。

不過在實際的案例中，我們發現有些使用者購買了 RDMA 的環境，但實際上並沒有用起來 GDR 技術，導致通訊效率不高。這裡列出了幾個關鍵的配置項，如果有類似問題的話，可以依序進行檢查和設定。

前面介紹了當前的一些主要的效能最佳化想法和方案，整體來看，無論I/O 最佳化、運算最佳化、通訊最佳化，最樸素的最佳化想法主要就是如何優化操作本身、或是否可以減少操作發生的次數，或者是否可以操作和其他的過程並行起來從而隱藏開銷等。

3. AIAK-Tranining 加速套件實踐

前面介紹了很多的優化工作，要想正確的使能，需要每個用戶對於框架的工程實現原理比較清晰。為了簡化訓練優化的成本，我們建造了 AIAK-Training 這個加速套件。

AIAK-Training 會圍繞數據加載、模型計算、通信等方面，構建全鏈路優化能力，同時我們會把這種優化能力，封裝成簡單易用的接口，用戶插入幾行代碼，即可比較方便的整合使用。同時，我們也正在建立自動化策略組合調優的機制，自動幫助使用者選擇有效的最佳化策略。

具體使用時，加速庫元件可以獨立的安裝部署，也可以直接使用我們提供的容器映像。

以下是一些具體的應用案例。

下圖所示，主要是針對 dataloader 的最佳化。在這個場景中，模型比較小，資料集規模也比較小，純計算的速度其實比較快，但是跨 EPOCH 的資料載入時間比較長，導致 I/O 耗時成為了主要的瓶頸。

透過使用 AIAK 中提供的流程重複使用、以及充分預取機制，整個模式訓練加速了  166%。

下圖是一個針對模型計算最佳化的案例。

Transformer 類別模型訓練場景，實際訓練時通訊擴展性接近線性，I/O 耗時佔比也非常低，計算是主要的效能瓶頸。

針對這個模型，AIAK-Training 進行了一系列計算層面的最佳化，包括主要結構的算子融合、混合精度、大 batch 調優等等，整個模型的訓練效率提升了 169%。

下圖的案例主要是應用了通訊層面的最佳化，在雲端上TCP 環境中啟用針對低頻寬網路的最佳化策略，在一些經典模型上例如resnet50、bert、vgg16，可以加速 26%~78%。

在自動駕駛場景中，我們也針對典型的2D 視覺、3D 視覺、雷射雷達，以及前融合類別的模型，做了一系列的模型訓練性能優化，訓練表現加速 49%~ 391%。

以上是AI 訓練加速原理解析與工程實務分享的詳細內容。更多資訊請關注PHP中文網其他相關文章！