這次分享將系統性的分析在 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 作業等彈性調度的需求。云原生 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 的訪問，從而節省訪存時間?；蛘咦龊糜嬎阒噶詈驮L存指令的 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%。