這次分享將系統(tǒng)性的分析在 AI 模型訓(xùn)練過(guò)程中的主要性能瓶頸，以及當(dāng)前針對(duì)這些瓶頸的主要的加速方案和技術(shù)原理，并介紹百度智能云在這方面的一些實(shí)踐成果。

今天的分享，主要包括三個(gè)部分：

首先介紹我們?yōu)槭裁葱枰?nbsp;AI 訓(xùn)練加速，也就是整體背景和出發(fā)點(diǎn)是什么；

第二部分我們會(huì)系統(tǒng)性的分析實(shí)際訓(xùn)練過(guò)程中的可能會(huì)遇到的性能瓶頸問(wèn)題，然后針對(duì)這些問(wèn)題，介紹目前主要的加速方案；

第三部分介紹百度百舸平臺(tái)的 AI 訓(xùn)練加速套件 AIAK-Training 在一些模型訓(xùn)練加速上的實(shí)踐效果。

1. 為什么需要 AI 訓(xùn)練加速？

在 AI 系統(tǒng)中，一個(gè)模型從生產(chǎn)到應(yīng)用，一般包括離線訓(xùn)練和推理部署兩大階段。

離線訓(xùn)練階段，就是產(chǎn)生模型的過(guò)程，用戶需要根據(jù)自己的任務(wù)場(chǎng)景，準(zhǔn)備好訓(xùn)練模型所需要的數(shù)據(jù)集，以及神經(jīng)網(wǎng)絡(luò)算法。

算法可以理解為是一個(gè)高度復(fù)雜的非凸數(shù)學(xué)函數(shù)，函數(shù)中包括很多變量以及參數(shù)。模型訓(xùn)練的過(guò)程其實(shí)就是在學(xué)習(xí)神經(jīng)網(wǎng)絡(luò)模型中的參數(shù)。

模型訓(xùn)練開(kāi)始后，會(huì)讀取數(shù)據(jù)，然后送入模型進(jìn)行前向計(jì)算，并計(jì)算與真實(shí)值的誤差。然后執(zhí)行反向計(jì)算得到參數(shù)梯度，最后更新參數(shù)。訓(xùn)練會(huì)進(jìn)行多輪的數(shù)據(jù)迭代。

訓(xùn)練完成之后，我們會(huì)保存訓(xùn)練好的模型，然后將模型做上線部署，接受用戶的真實(shí)輸入，通過(guò)前向計(jì)算，完成推理。

因此，無(wú)論是訓(xùn)練還是推理，核心都是數(shù)據(jù)計(jì)算。 為了加速計(jì)算效率，一般都是通過(guò) GPU 等異構(gòu)加速芯片來(lái)進(jìn)行訓(xùn)練和推理。

另外，從深度學(xué)習(xí)模型發(fā)展歷程來(lái)看，為了能夠持續(xù)突破模型的精度上限，模型參數(shù)量其實(shí)在快速的膨脹。然而更大的參數(shù)量，就會(huì)帶來(lái)更大的計(jì)算復(fù)雜度。

下圖左側(cè)是摘自一篇公開(kāi)的論文，從這篇總結(jié)里，我們看到在 2010 年之前，模型的計(jì)算量大約 20 個(gè)月翻一番。在 2010~2015 年，常規(guī)模型計(jì)算每 5-6 個(gè)月翻一番。而在 2015 年之后，衍生了大模型訓(xùn)練的趨勢(shì)，計(jì)算量增長(zhǎng) 10~100 倍。

模型訓(xùn)練對(duì)算力以及基礎(chǔ)設(shè)施的要求越來(lái)越高，訓(xùn)練需要更多的算力，也需要更長(zhǎng)的時(shí)間，這也導(dǎo)致了需要更多的資源成本。這里我們列舉了一些論文或研究中公開(kāi)的成本數(shù)據(jù)，反應(yīng)了模型訓(xùn)練的費(fèi)用是非常高昂的。

因此，如何穩(wěn)定的進(jìn)行模型訓(xùn)練，如何持續(xù)降本增效其實(shí)至關(guān)重要。

在這樣的大背景下，百度智能云推出了百度百舸 · AI 異構(gòu)計(jì)算平臺(tái)，目標(biāo)是為 AI 場(chǎng)景提供軟硬一體化的解決方案。通過(guò) AI 計(jì)算、AI 存儲(chǔ)、AI 加速、AI 容器四層技術(shù)棧，滿足上層業(yè)務(wù)場(chǎng)景的需求。

• AI 計(jì)算層，提供了包括高性能的 GPU、以及昆侖等異構(gòu)芯片資源，以及高性能的 RDMA 或 IB 網(wǎng)絡(luò)，以及自研的超級(jí) AI 計(jì)算機(jī) X-MAN 等；
• AI 存儲(chǔ)層，包括對(duì)象存儲(chǔ) BOS 滿足數(shù)據(jù)湖存儲(chǔ)的需求、以及專為 AI 設(shè)計(jì)的高性能并行文件系統(tǒng) PFS；
• AI 加速層，包括數(shù)據(jù)湖存儲(chǔ)加速套件 RapidFS，AI 訓(xùn)練加速套件 AIAK-Training，AI 推理加速套件 AIAK-Inference；
• AI 容器層，也即是資源調(diào)度層，利用云原生的技術(shù)能力，滿足 GPU、AI 作業(yè)等彈性調(diào)度的需求。云原生 AI 的內(nèi)容在我們上一期的技術(shù)公開(kāi)課有專門分享。

當(dāng)我們考慮做性能加速的時(shí)候，第一個(gè)想到的可能是使用更好的硬件。

這會(huì)帶來(lái)一定的性能提升，但是大部分情況下可能并沒(méi)有充分發(fā)揮出硬件的計(jì)算能力，核心的原因就是訓(xùn)練代碼的執(zhí)行效率并沒(méi)有調(diào)到最優(yōu)或更優(yōu)的狀態(tài)。

• 首先，框架為了兼容更多的需求，一般會(huì)提供一些較為通用的優(yōu)化能力，更多的會(huì)關(guān)注易用性、生態(tài)建設(shè)等；
• 其次，算法工程師在設(shè)計(jì)算法時(shí)，核心的精力在如何提高模型的精度，性能效率關(guān)注不足，某些情況下框架提供的一些通用優(yōu)化能力，也都沒(méi)有利用起來(lái)。

因此，當(dāng)我們決定使用某種模型算法時(shí)，為了達(dá)到較好的資源效率和訓(xùn)練效率，我們需要有意識(shí)的去優(yōu)化。不過(guò)這里也有很多的技術(shù)挑戰(zhàn)：

• 性能影響因素比較多，需要根據(jù)不同模型自身的運(yùn)行特點(diǎn)進(jìn)行分析，沒(méi)有完全固定的優(yōu)化方案；
• 性能優(yōu)化需要理解工程實(shí)現(xiàn)的原理，比如當(dāng)我們?nèi)プ霎悩?gòu)芯片計(jì)算的優(yōu)化，需要專業(yè)的異構(gòu)研發(fā)經(jīng)驗(yàn)才能開(kāi)展，技術(shù)門檻較高；
• 還有些情況，當(dāng)我們做分布式訓(xùn)練時(shí)，可能還需要結(jié)合集群硬件和網(wǎng)絡(luò)拓?fù)?，?lái)優(yōu)化分布式訓(xùn)練場(chǎng)景下的擴(kuò)展性問(wèn)題。

這些挑戰(zhàn)極大地影響了模型訓(xùn)練性能的調(diào)優(yōu)，因此我們推出了AIAK-Training 的加速套件，期望通過(guò)抽象易用性的接口降低優(yōu)化成本，并通過(guò)軟硬協(xié)同的優(yōu)化手段，來(lái)充分加速客戶在百度智能云上的模型訓(xùn)練性能。

2. 訓(xùn)練性能開(kāi)銷分析和加速方案

在介紹 AIAK-Training 具體效果之前，我們先介紹下訓(xùn)練加速這個(gè)話題下關(guān)鍵的技術(shù)思路和方案原理是什么樣的。

因?yàn)槟Ｐ陀?xùn)練優(yōu)化本身是一個(gè)軟硬件綜合的工作，技術(shù)棧相對(duì)比較復(fù)雜，今天的內(nèi)容肯定沒(méi)辦法涵蓋全部細(xì)節(jié)，我們盡量把關(guān)鍵思路講到。

首先我們看下當(dāng)前的模型訓(xùn)練方案。過(guò)去的發(fā)展階段里，模型訓(xùn)練方案關(guān)鍵有兩個(gè)層次的變化，一是從單卡訓(xùn)練到分布式訓(xùn)練的變化，二是從數(shù)據(jù)并行訓(xùn)練到多維混合并行訓(xùn)練的變化。這里的核心驅(qū)動(dòng)點(diǎn)一個(gè)是訓(xùn)練的數(shù)據(jù)量，一個(gè)是模型的參數(shù)量。

• 單卡訓(xùn)練方式：實(shí)際采用這種模式，一般都是模型參數(shù)量和數(shù)據(jù)量相對(duì)比較少，單卡的訓(xùn)練時(shí)間可以接受。模型參數(shù)規(guī)模需要保證在訓(xùn)練過(guò)程中，單張卡的顯存能夠滿足存儲(chǔ)的上限。按照新的 GPU 卡的顯存容量配置，一般可以放置的最大規(guī)模在10 億級(jí)參數(shù)量；
• 當(dāng)數(shù)據(jù)集規(guī)模比較大的時(shí)候，因?yàn)橛?xùn)練過(guò)程中需要多次遍歷全量的數(shù)據(jù)集，單卡訓(xùn)練就會(huì)需要比較長(zhǎng)的時(shí)間，這時(shí)可以從單卡擴(kuò)展到多卡，通過(guò)分布式的方式來(lái)加快訓(xùn)練。

這里應(yīng)用最廣泛的就是數(shù)據(jù)并行。在數(shù)據(jù)并行方案下，數(shù)據(jù)集會(huì)被平均切分成多份，然后每張卡上保存完整的模型，各自獨(dú)立并行處理切分后的子數(shù)據(jù)集。

當(dāng)模型參數(shù)量足夠大的時(shí)候，比如參數(shù)量達(dá)到百億、千億級(jí)別，單卡放不下完整的模型，這里又出現(xiàn)了模型并行、或者同時(shí)使用數(shù)據(jù)并行和模型并行的混合并行方案。

模型并行，會(huì)將模型切分到不同的卡上，每個(gè)卡上放置模型的一部分，這里又根據(jù)切分的方式不同，比如層內(nèi)切分或?qū)娱g切分，又細(xì)分了 Tensor 并行和流水線并行的方式。

因?yàn)樵谝话隳Ｐ陀?xùn)練中，使用更多的還是數(shù)據(jù)并行，我們下面重點(diǎn)還是以數(shù)據(jù)并行為例，來(lái)介紹性能優(yōu)化的思路。

我們從軟硬件整體視角先理解下單卡訓(xùn)練過(guò)程中存在的性能開(kāi)銷。

下圖左邊是我們從軟件角度上的訓(xùn)練流程。單卡訓(xùn)練的過(guò)程，主要包括數(shù)據(jù)讀取、數(shù)據(jù)預(yù)處理，前向計(jì)算輸出并計(jì)算 loss，根據(jù) loss 函數(shù)反向計(jì)算，得到每一層參數(shù)的梯度，最后根據(jù)梯度更新模型參數(shù)。持續(xù)該過(guò)程，直到訓(xùn)練收斂。

下圖右邊，是一個(gè)簡(jiǎn)化的節(jié)點(diǎn)硬件拓?fù)鋱D。最上面是數(shù)據(jù)存儲(chǔ)，可以是本地存儲(chǔ)，也可以是網(wǎng)絡(luò)存儲(chǔ)。然后是 CPU、內(nèi)存，CPU 下通過(guò)多個(gè) PCIe Switch 連接著 8 張 GPU 卡，編號(hào)從 0~7，8 張卡之間通過(guò) NVSwitch 互聯(lián)。不同的計(jì)算實(shí)例，硬件的拓?fù)浣Y(jié)構(gòu)會(huì)有不同。

• 當(dāng)訓(xùn)練啟動(dòng)時(shí)，首先數(shù)據(jù)讀取，涉及從存儲(chǔ)介質(zhì)中讀數(shù)據(jù)到內(nèi)存中，主要是存儲(chǔ)I/O 開(kāi)銷。根據(jù)存儲(chǔ)介質(zhì)，要讀取的數(shù)據(jù)量不同，這部分時(shí)間開(kāi)銷也不同；
• 接下來(lái)是數(shù)據(jù)預(yù)處理，主要是對(duì)讀入的數(shù)據(jù)進(jìn)行一些數(shù)據(jù)增強(qiáng)的操作，比如對(duì)圖片進(jìn)行 resize，對(duì)比度、飽和度調(diào)整等等，這部分工作大多數(shù)情況下是在 CPU 上。當(dāng)在 CPU 上完成數(shù)據(jù)預(yù)處理操作之后，需要從主機(jī)內(nèi)存拷貝到 GPU 顯存上，涉及到主機(jī)和設(shè)備之間的內(nèi)存拷貝開(kāi)銷；
• 然后開(kāi)始前向計(jì)算、反向計(jì)算、參數(shù)更新，這部分更多操作主要是通過(guò) GPU 來(lái)進(jìn)行，主要時(shí)間花費(fèi)在 GPU 計(jì)算上面。這個(gè)過(guò)程里也可能會(huì)穿插一些 CPU 上的操作，可能也需要做主機(jī)和設(shè)備內(nèi)存之間的拷貝。

因此，從單卡角度看，主要存在 I/O、CPU 預(yù)處理、CPU 和 GPU 之間數(shù)據(jù)拷貝，GPU 計(jì)算等方面的開(kāi)銷。

接著我們看下數(shù)據(jù)并行的過(guò)程。

下圖左邊還是訓(xùn)練的主流程，右邊展示了一個(gè) 3 機(jī) 24 卡的訓(xùn)練集群的硬件拓?fù)洌? 臺(tái)機(jī)器通過(guò)網(wǎng)絡(luò)互聯(lián)。

前面的部分我們也介紹到了，在數(shù)據(jù)并行里每個(gè)設(shè)備并行獨(dú)立地執(zhí)行前向和反向計(jì)算過(guò)程，因此，每個(gè)訓(xùn)練進(jìn)程也都會(huì)遇到前面講的單卡訓(xùn)練中的性能開(kāi)銷問(wèn)題。

數(shù)據(jù)并行為了保證和單卡訓(xùn)練在數(shù)學(xué)上等價(jià)，需要確保每張卡的模型參數(shù)在迭代過(guò)程中始終保持一致。這里一方面需要讓各 GPU 卡的模型參數(shù)初始化狀態(tài)一致，這個(gè)一般是在訓(xùn)練開(kāi)始前，通過(guò)廣播的方式將第一張卡上的參數(shù)狀態(tài)廣播到其他的卡。

而在訓(xùn)練期間，由于每個(gè)設(shè)備處理的數(shù)據(jù)不同，前向計(jì)算所得到的模型損失值也是不同的，因此還需要在每個(gè)設(shè)備反向計(jì)算出梯度之后，進(jìn)行梯度的平均，以平均后的梯度值來(lái)更新模型參數(shù)，從而保證每張卡的模型參數(shù)在迭代過(guò)程中始終保持一致。

梯度平均涉及到通信的過(guò)程，包括節(jié)點(diǎn)內(nèi)部卡之間的通信，以及跨節(jié)點(diǎn)的網(wǎng)絡(luò)通信開(kāi)銷。這里的通信又包括同步通信和異步通信，不過(guò)為了保證模型訓(xùn)練的收斂，一般都是采用同步通信的方案，后續(xù)的優(yōu)化工作也都是基于同步通信的方式來(lái)展開(kāi)。

由上可知，數(shù)據(jù)并行相比單卡訓(xùn)練，主要增加了額外的通信開(kāi)銷。

通過(guò)前述分析，我們知道加速AI 訓(xùn)練不單是某一方面的工作，需要從數(shù)據(jù)加載、模型計(jì)算、分布式通信等系統(tǒng)維度綜合考慮。這里說(shuō)的數(shù)據(jù)加載，包括數(shù)據(jù) I/O、預(yù)處理、內(nèi)存拷貝等過(guò)程。

在具體的優(yōu)化實(shí)踐中，給定一個(gè)待優(yōu)化的模型，加速模型訓(xùn)練就是要不斷提升訓(xùn)練的總體吞吐（每秒可以訓(xùn)練的樣本數(shù)）。這個(gè)過(guò)程，我們一般可以先分析單卡的訓(xùn)練吞吐，當(dāng)單卡的訓(xùn)練吞吐提升上來(lái)之后，我們?cè)倏磸膯慰〝U(kuò)展到多卡，看看如何提升多卡的訓(xùn)練加速比。

首先單卡的訓(xùn)練優(yōu)化，極限優(yōu)化目標(biāo)是全部時(shí)間都在 GPU 計(jì)算上，加速器利用率100%。當(dāng)然實(shí)際很難完全達(dá)到這個(gè)狀態(tài)，但是我們可以按這個(gè)指標(biāo)來(lái)牽引或衡量我們的工作。

單卡性能優(yōu)化關(guān)鍵包括兩部分：

• 首先數(shù)據(jù)加載效率的優(yōu)化。從存儲(chǔ)系統(tǒng)上，我們可以使用更高性能的存儲(chǔ)介質(zhì)，或者基于這些高速存儲(chǔ)介質(zhì)組成的并行文件系統(tǒng)，或者說(shuō)一些緩存加速系統(tǒng)。前面介紹到的，百度百舸也提供了相應(yīng)的存儲(chǔ)系統(tǒng)方案，比如 PFS、RapidFS 等。除此之外，還需要在框架 dataloader 中優(yōu)化數(shù)據(jù)的讀取過(guò)程。
• 其次就是模型計(jì)算效率的優(yōu)化。主要考慮如何優(yōu)化計(jì)算實(shí)現(xiàn)，怎么提升計(jì)算單元的利用效率，這塊可能就需要結(jié)合模型具體分析。

然后從單卡擴(kuò)展到多卡，目標(biāo)是如何達(dá)到線性加速比。線性加速比這個(gè)指標(biāo)，簡(jiǎn)單來(lái)說(shuō)就是從 1 張卡擴(kuò)到 2 張卡訓(xùn)練時(shí)，訓(xùn)練的性能是否是單卡的2倍。

這里核心就是優(yōu)化分布式的通信效率，一方面是硬件層面的優(yōu)化，另外一方面在實(shí)際通信中，需要考慮怎么利用好網(wǎng)絡(luò)的帶寬資源，或者是否能夠?qū)⑼ㄐ胚^(guò)程進(jìn)行隱藏等。

下面我們分別從這幾個(gè)方面詳細(xì)展開(kāi)。

首先是數(shù)據(jù)加載方面的優(yōu)化。

當(dāng)我們實(shí)例化一個(gè) dataloader 之后，我們會(huì)持續(xù)迭代 dataloader 來(lái)讀取一個(gè) batch 的數(shù)據(jù)進(jìn)行模型訓(xùn)練。

如果這里不做任何優(yōu)化，如下圖上半部分所示，每個(gè) batch 的數(shù)據(jù)加載過(guò)程和每個(gè) batch 的模型訓(xùn)練過(guò)程，實(shí)際上是串行進(jìn)行的。從 GPU 視角來(lái)看，就會(huì)出現(xiàn)因?yàn)閿?shù)據(jù)加載導(dǎo)致的計(jì)算間隙，計(jì)算資源存在時(shí)間上的浪費(fèi)。

除了前面講到我們用更好的硬件直接提升數(shù)據(jù)讀的效率之外，還可以怎么優(yōu)化呢？

實(shí)際在 AI 訓(xùn)練過(guò)程中，數(shù)據(jù)訪問(wèn)上有兩個(gè)關(guān)鍵特征：

• 當(dāng)數(shù)據(jù)集做完 shuffle 之后，每輪訓(xùn)練所需要的 batch 數(shù)據(jù)以及訪問(wèn)順序是已知的；
• 任意兩個(gè) batch 的數(shù)據(jù)讀可以并行，因?yàn)閿?shù)據(jù)之間沒(méi)有任何依賴關(guān)系。

因此，我們?cè)诓灰蕾囉布用娓膭?dòng)的時(shí)候，可以做的一個(gè)優(yōu)化工作就是數(shù)據(jù)預(yù)取，當(dāng)訓(xùn)練第一個(gè) batch 數(shù)據(jù)的時(shí)候，可以提前加載下一個(gè) batch 數(shù)據(jù)，讓 I/O 的過(guò)程和 GPU 上的計(jì)算充分并行起來(lái)。

首先，我們需要利用好 dataloader 中已有的優(yōu)化方案，一是合理設(shè)置 num_workers 超參數(shù)，通過(guò)多進(jìn)程的方式讀數(shù)據(jù)，這一步可以實(shí)現(xiàn)數(shù)據(jù)從存儲(chǔ)系統(tǒng)中預(yù)取到主機(jī)內(nèi)存。二是從主機(jī)內(nèi)存拷貝到 GPU顯存，可以通過(guò) pinned memory 的機(jī)制來(lái)加速。

大概介紹下 pinned memory 加速的主要原理：內(nèi)存數(shù)據(jù)有兩種類型 pageable memory 和 pinned memory，pageable memory 中的數(shù)據(jù)有被換出到磁盤上的可能。這種情況下，當(dāng)執(zhí)行 H2D 時(shí)，可能需要先從磁盤讀到內(nèi)存，然后從內(nèi)存拷貝到顯存。另外，pageable memory 數(shù)據(jù)拷貝到 GPU 顯存時(shí)，需要先創(chuàng)建一個(gè)臨時(shí)的 pinned memory 緩沖區(qū)，把數(shù)據(jù)從 pageable memory 拷貝 pinned memory，之后才能傳輸?shù)?nbsp;GPU 上，也多了額外的數(shù)據(jù)搬運(yùn)的操作。

不過(guò)當(dāng)我們使能了上述方案后，我們僅實(shí)現(xiàn)了從存儲(chǔ)系統(tǒng)到主機(jī)內(nèi)存的預(yù)取，加快了主機(jī)到設(shè)備的數(shù)據(jù)拷貝速度。但是主機(jī)到設(shè)備的內(nèi)存拷貝，和實(shí)際計(jì)算 kernel 在 GPU 上還是串行執(zhí)行，也就是 GPU 上依然存在少量的時(shí)間間隙。

AIAK 針對(duì)這個(gè)問(wèn)題，又做了進(jìn)一步的優(yōu)化，可以實(shí)現(xiàn) H2D 和前向計(jì)算的 overlap。

在數(shù)據(jù)并行場(chǎng)景下，還需要注意的一個(gè)事情，數(shù)據(jù)需要均衡的切分。

如果每個(gè)訓(xùn)練進(jìn)程分配的數(shù)據(jù)不均衡，計(jì)算量就會(huì)不同，也就導(dǎo)致每個(gè)進(jìn)程前向計(jì)算和反向計(jì)算完成的時(shí)間不同，那么先完成計(jì)算的進(jìn)程，在反向過(guò)程中就會(huì)先進(jìn)入到梯度通信環(huán)節(jié)中，但是因?yàn)?Allreduce 通信是同步通信操作，需要所有進(jìn)程同時(shí)開(kāi)始并同時(shí)結(jié)束，因此先開(kāi)始通信的進(jìn)程，會(huì)一直等待其他所有進(jìn)程也發(fā)起了 AllReduce 后才能一起完成通信操作。這里就會(huì)出現(xiàn)因?yàn)榭炻灰粚?dǎo)致的資源空閑問(wèn)題。

為了解決這種問(wèn)題，需要每個(gè)進(jìn)程使用相同的 batchsize 來(lái)讀取數(shù)據(jù)，另外每個(gè) batch 的數(shù)據(jù)量要均衡。圖像類數(shù)據(jù)一般會(huì)固定尺寸進(jìn)行訓(xùn)練，而像 NLP 類模型需要處理變長(zhǎng)的語(yǔ)句，可能需要進(jìn)行特殊的處理，比如可以將數(shù)據(jù) padding 到相同的長(zhǎng)度，或者通過(guò)樣本長(zhǎng)度排序的方式來(lái)均衡分配等。

下面介紹計(jì)算效率的優(yōu)化。

計(jì)算包括前向、反向、參數(shù)更新。優(yōu)化計(jì)算的目標(biāo)，是為了能夠充分發(fā)揮出異構(gòu)硬件的算力，理想情況就是讓 GPU 芯片實(shí)際計(jì)算時(shí)的性能達(dá)到理論峰值。

我們先從一個(gè)單算子的角度分析，當(dāng)我們準(zhǔn)備在 GPU 上執(zhí)行一個(gè)計(jì)算操作的時(shí)候，簡(jiǎn)化的流程有四步。

• 首先在 CPU 上異步發(fā)射一個(gè) GPU 計(jì)算 Kernel；
• 當(dāng) Kernel 調(diào)度執(zhí)行時(shí)，需要先從 GPU 上的 Global Memory 讀取計(jì)算所需要的數(shù)據(jù);
• 開(kāi)始執(zhí)行計(jì)算;
• 當(dāng)計(jì)算完成后，需要將計(jì)算的結(jié)果寫回 Global Memory。

根據(jù)計(jì)算、訪存的開(kāi)銷占比，一般會(huì)將算子分類為計(jì)算瓶頸或者訪存瓶頸。

當(dāng)從一個(gè)算子擴(kuò)展到一個(gè)完整的模型訓(xùn)練時(shí)，因?yàn)橐B續(xù)執(zhí)行非常多的計(jì)算 Kernel，那么 Kernel 計(jì)算之間就會(huì)出現(xiàn)很多因?yàn)?nbsp;Kernel Launch、中間結(jié)果的讀寫導(dǎo)致的計(jì)算間隙問(wèn)題。

由上可知，優(yōu)化模型計(jì)算效率，需要從訪存優(yōu)化、計(jì)算優(yōu)化、其他開(kāi)銷的優(yōu)化綜合考慮。

• 訪存優(yōu)化，主要考慮如何減少數(shù)據(jù)在顯存和計(jì)算單元之間搬運(yùn)的時(shí)間。從算子實(shí)現(xiàn)角度上，需要利用好 GPU 存儲(chǔ)層次架構(gòu)，比如把數(shù)據(jù)搬運(yùn)到更快的存儲(chǔ)器上比如 share memory，減少對(duì) global memory 的訪問(wèn)，從而節(jié)省訪存時(shí)間。或者做好計(jì)算指令和訪存指令的 overlap，來(lái)提升計(jì)算訪存比。而從單算子擴(kuò)展到多算子時(shí)，還需要考慮如何減少中間結(jié)果的讀寫，這種問(wèn)題一般可以通過(guò)算子融合的手段來(lái)優(yōu)化；
• 計(jì)算優(yōu)化，計(jì)算瓶頸問(wèn)題多半應(yīng)該是沒(méi)有正確對(duì)任務(wù)進(jìn)行分塊，或者沒(méi)有利用好 GPU 并行計(jì)算的優(yōu)勢(shì)，導(dǎo)致并行度不高。還有可能沒(méi)有使用合并指令集導(dǎo)致計(jì)算效率低下，或者沒(méi)有使用 Tensor Core 等高性能計(jì)算單元，造成資源浪費(fèi)。不同的問(wèn)題，也對(duì)應(yīng)著不同的優(yōu)化手段；
• Kernel Launch 等其他開(kāi)銷，大量時(shí)間花費(fèi)在訪存或計(jì)算之外，可以考慮的優(yōu)化手段如算子融合、Cuda Graph 等。

首先是算子融合。算子在底層 GPU 執(zhí)行時(shí)，會(huì)發(fā)起一次或者多次 Kernel Launch，Kernel 之間交互數(shù)據(jù)也需要經(jīng)過(guò)顯存，而算子融合就是將多個(gè) GPU Kernel 融合成一個(gè)大 Kernel，統(tǒng)一發(fā)起和執(zhí)行。

• 因?yàn)闇p少了需要執(zhí)行的算子數(shù)量，從而可以減少 Kernel 調(diào)度和發(fā)起的開(kāi)銷；
• 通過(guò)融合，可以通過(guò)寄存器等來(lái)傳遞中間結(jié)果，避免從 global memory 的來(lái)回搬運(yùn)，極大降低了顯存等待的時(shí)間；
• 在某些場(chǎng)景中，可以通過(guò)算子融合，可以更充分的利用計(jì)算資源，提升資源和計(jì)算的效率。

算子融合具體如何實(shí)現(xiàn)呢？

一種方式，分析模型中的低效操作，專家經(jīng)驗(yàn)手寫融合算子。在 GPU 上主要就是 CUDA 算子研發(fā)，這里存在一定的門檻。AIAK-Training 會(huì)針對(duì)典型的模型結(jié)構(gòu)，或者客戶需求，提供高效優(yōu)化的算子實(shí)現(xiàn)。

另一種方式，就是編譯優(yōu)化的方案。通過(guò)編譯的方式進(jìn)行計(jì)算優(yōu)化，以及代碼自動(dòng)生成，從而降低在不同硬件上的手工優(yōu)化成本。不過(guò)當(dāng)前很多編譯方案更多還是針對(duì)推理優(yōu)化，訓(xùn)練上的方案還在快速演進(jìn)過(guò)程中。不過(guò)從極致性能角度看，未來(lái)一段時(shí)間依然離不開(kāi)手寫融合算子的工作。

下面介紹一些算子融合的實(shí)際案例。第一個(gè)是針對(duì)典型模型網(wǎng)絡(luò)結(jié)構(gòu)的優(yōu)化。

下圖展示了我們?cè)?SwinTransformer 模型中針對(duì)核心模塊 WindowAttention 進(jìn)行的計(jì)算融合優(yōu)化。

WindowAttention 結(jié)構(gòu)，核心操作公式如下圖所示。計(jì)算過(guò)程中，需要依次執(zhí)行 7 個(gè)計(jì)算 Kernel。再加上一些 reshape 等轉(zhuǎn)換操作，總共需要 launch 10 個(gè) Kernel。通過(guò)性能分析發(fā)現(xiàn)，實(shí)際執(zhí)行過(guò)程中 launch kernel 的間隔冗余開(kāi)銷占到了端到端 80% 以上的時(shí)間，導(dǎo)致該模塊存在著較大的優(yōu)化空間。

通過(guò)將這些 Kernel 融合成一個(gè)，整個(gè)模塊的執(zhí)行時(shí)間從 392 微秒減少到 13 微秒，單算子加速了 30 倍。整個(gè)模型的訓(xùn)練效率，端到端加速了 20% 以上。

這個(gè)優(yōu)化的核心思路，主要有三點(diǎn)：

• 是利用好 GPU 的三級(jí)訪存流水：顯存、share memory、寄存器；
• 通過(guò)分塊策略，將 2 個(gè)矩陣乘法和 softmax 進(jìn)行融合；
• 針對(duì)前向中間結(jié)果的優(yōu)化，充分利用 Tensor Core，在反向計(jì)算中使用重計(jì)算代替重加載，使得訪存開(kāi)銷極大降低。

下圖是一個(gè)數(shù)據(jù)操作的融合舉例，是在 FCOS3D 模型中對(duì)于坐標(biāo)壓縮操作的一個(gè)優(yōu)化。

通過(guò)性能分析發(fā)現(xiàn)，這個(gè)操作過(guò)程中存在大量 GPU 空隙，GPU 利用率較低。該操作主要的功能是根據(jù) index 將 3D-Tensor 壓縮成 2D-Tensor。在原生實(shí)現(xiàn)中會(huì)先在host 端生成 index，再進(jìn)行 H2D 拷貝，最后完成 Tensor 壓縮，這會(huì)造成額外的拷貝和等待開(kāi)銷。

為此，我們重新實(shí)現(xiàn)這部分操作，核心思路主要就是將操作全部遷移至 GPU 上，直接在 GPU 上完成 index 的生成和 Tensor 的壓縮，減少 CPU 參與，同時(shí)避免了非必要的 CPU-GPU 之間的內(nèi)存拷貝。

單算子執(zhí)行時(shí)間從 9.69 毫秒減少到 32 微秒，加速了 300 倍，整個(gè)模型端到端訓(xùn)練提升了 10% 以上。

下面我們介紹另一個(gè)計(jì)算優(yōu)化的思路，就是提高計(jì)算的并行度，充分利用 GPU 并行計(jì)算的優(yōu)勢(shì)，同樣借助一些實(shí)際案例來(lái)介紹。

我們發(fā)現(xiàn)在一些模型中，有一些操作是串行執(zhí)行的。比如在一些目標(biāo)檢測(cè)模型里，在 loss 計(jì)算過(guò)程中，有些操作并不是按照一個(gè) batch 進(jìn)行操作，而是 for-loop 每張圖片或一個(gè)樣本，這種情況下，當(dāng)我們?nèi)ヌ嵘?nbsp;batchsize 的時(shí)候，因?yàn)檫@里的串行，可能沒(méi)法達(dá)到我們想要的性能效果。

以 YOLOv7 中的 SimOTA 操作舉例，原生實(shí)現(xiàn)中，是通過(guò) for-loop 遍歷一個(gè) batch 的每一張圖片，然后為圖片的 gtbox 執(zhí)行 SimOTA 標(biāo)簽分配。這種串行的實(shí)現(xiàn)方式導(dǎo)致該部分操作的 GPU 利用率非常低效。

而每張圖片處理時(shí)數(shù)據(jù)之間是沒(méi)有依賴的。因此，我們做的一項(xiàng)工作就是將串行計(jì)算改成 batch 并行計(jì)算，通過(guò)對(duì)一個(gè) batch 的數(shù)據(jù)進(jìn)行并行化的標(biāo)簽分配，從而加速這部分計(jì)算的效率。

最終效果上，SimOTA 操作的耗時(shí)從 384 毫秒下降到 69 毫秒，計(jì)算效率提升 5.5 倍，整個(gè)模型的端到端訓(xùn)練效率提升了 18% 以上。

模型訓(xùn)練過(guò)程中也有其他的類似場(chǎng)景，比如說(shuō)參數(shù)更新。參數(shù)更新時(shí)，默認(rèn)也是通過(guò)循環(huán)的方式，遍歷每個(gè)參數(shù)，然后每個(gè)參數(shù)都會(huì)啟動(dòng)一個(gè)參數(shù)更新的 Cuda Kernel，然后依次執(zhí)行。

針對(duì)這種情況，AIAK 也增加了 FusedOptimizer 的優(yōu)化，通過(guò)融合參數(shù)更新的算子，實(shí)現(xiàn)批量化更新參數(shù)，大幅減少 Kernel Launch 次數(shù)。

下面介紹了另一個(gè)優(yōu)化手段 CUDA Graph，主要是為了減少 CPU Launch Kernel 的開(kāi)銷。

CUDA Graph 是在 CUDA 10 版本中引入的特性，可以將一系列 CUDA Kernel 封裝成單個(gè)單元，可以通過(guò)一次 CPU Launch 操作來(lái)啟動(dòng)多個(gè) GPU Kernel，從而減少了CPU Launch Kernel 的開(kāi)銷。

如下圖所示，默認(rèn)情況下 CPU 需要依次發(fā)射多個(gè) Kernel，如果 Kernel 計(jì)算時(shí)間比較短，Kernel 之間的 Launch 間隙可能成為性能瓶頸。通過(guò) CUDA Graph，僅需要花費(fèi)一些額外時(shí)間構(gòu)建 Graph，后續(xù)通過(guò)一次 Graph 的發(fā)射，即可大幅縮短實(shí)際執(zhí)行時(shí) Kernel 之間的間隙。

現(xiàn)在很多框架也都增加了對(duì) CUDA Graph 的支持，通過(guò)插入一些代碼來(lái)使K能這個(gè)功能。不過(guò)也有一些使用上的限制，比如不支持動(dòng)態(tài) shape、不支持動(dòng)態(tài)控制流、過(guò)程中不能捕獲 CPU 操作等等，可以根據(jù)模型情況，嘗試使用下這種優(yōu)化能力。

下面介紹最后一個(gè)計(jì)算優(yōu)化手段，充分利用 Tensor Core 計(jì)算單元。

一個(gè) GPU 內(nèi)一般包含多個(gè) SM，每個(gè) SM 包括不同數(shù)據(jù)類型的計(jì)算核心，以及各類存儲(chǔ)資源等。下圖左邊所示，是一個(gè) NVIDIA A100 SM 的示意圖，一個(gè) SM 包含 64 個(gè) FP32 CUDA Core，以及 4 個(gè) Tensor Core。

模型訓(xùn)練時(shí)，默認(rèn)情況主要是用到 FP32 CUDA Core 計(jì)算，而 Tensor Core 是一塊特殊的硬件執(zhí)行單元，是從 Volta 系列 GPU 開(kāi)始引入，主要是用來(lái)加速矩陣或卷積的操作效率。

相比 FP32 CUDA Core 一次只能執(zhí)行兩個(gè)標(biāo)量的計(jì)算，Tensor Core 可以一次對(duì)兩個(gè)矩陣執(zhí)行計(jì)算，因此 Tensor Core 的計(jì)算吞吐遠(yuǎn)高于 FP32 CUDA Core。

在 A100 中，Tensor Core 支持了多種浮點(diǎn)數(shù)據(jù)類型，對(duì)于深度學(xué)習(xí)訓(xùn)練來(lái)說(shuō)，可能涉及到包括 FP16、BF16、以及 TF32 模式。

TF32 主要用于單精度訓(xùn)練場(chǎng)景，相比 FP32 訓(xùn)練，在保持相同的訪存帶寬需求下，理論計(jì)算吞吐量提高了 8 倍。

FP16/BF16 主要用于混合精度訓(xùn)練場(chǎng)景，相比 FP32 訓(xùn)練，訪存需求減少了一半，理論計(jì)算吞吐量提高了 16 倍。

使用 Tensor Core 的方式，可以使用底層的 cublas 或 cuda 接口進(jìn)行編程，而對(duì)于算法開(kāi)發(fā)者來(lái)說(shuō)，更直接的就是使用框架中提供的 TF32 訓(xùn)練或混合精度訓(xùn)練方案。

首先是 TF32 訓(xùn)練模式，TF32 是 Ampere 開(kāi)始引入。

TF32 在浮點(diǎn)數(shù)的表達(dá)中，有 8 個(gè)指數(shù)位，10 個(gè)尾數(shù)位和 1 個(gè)符號(hào)位。指數(shù)位與  FP32 相同，即數(shù)據(jù)表示范圍相同，但是尾數(shù)位低于 FP32，和 FP16 相同。

需要注意的是，TF32 不是一個(gè)對(duì)外開(kāi)放的數(shù)值類型，只是 Tensor Core 的一種計(jì)算模式，也就是用戶不能去直接創(chuàng)建一個(gè) TF32 類型的浮點(diǎn)數(shù)。

當(dāng)使能 TF32 的時(shí)候，Tensor Core 計(jì)算矩陣或卷積操作時(shí)，會(huì)自動(dòng)將 FP32 轉(zhuǎn)換成 TF32，計(jì)算完成之后，輸出的數(shù)據(jù)類型依然是 FP32 類型。

TF32 訓(xùn)練在某些框架版本中是默認(rèn)開(kāi)啟，某些框架版本中可能需要通過(guò)環(huán)境變量或者參數(shù)配置來(lái)手工開(kāi)啟，具體需要參考框架的用戶手冊(cè)。

不過(guò)由于 TF32 相比 FP32 來(lái)說(shuō)，精度范圍降低了，實(shí)際訓(xùn)練時(shí)還需要關(guān)注對(duì)模型收斂精度的影響。

混合精度訓(xùn)練是指在盡可能減少模型精度損失的情況下，使用 FP32 和 FP16 混合精度進(jìn)行訓(xùn)練。

混合精度訓(xùn)練的收益主要有：相比 FP32 訓(xùn)練，內(nèi)存需求減少，可以訓(xùn)練更大的網(wǎng)絡(luò)或使用更大的 batchsize。使用更少的內(nèi)存帶寬，可以加速數(shù)據(jù)傳輸，半精度的計(jì)算也可以讓數(shù)學(xué)運(yùn)算效率更快；

不過(guò)因?yàn)?FP16 的指數(shù)位和尾數(shù)位的范圍都比 FP32 要少，因此數(shù)值表示范圍和精度都會(huì)有降低，在實(shí)際使用的時(shí)候，就可能出現(xiàn)因?yàn)楸硎痉秶M窄導(dǎo)致的數(shù)值溢出問(wèn)題，或者因?yàn)榫炔蛔銓?dǎo)致舍入誤差。

為了優(yōu)化類似問(wèn)題，混合精度訓(xùn)練方案中有幾個(gè)關(guān)鍵的技術(shù)工作：

• 算子黑白名單機(jī)制，框架使用黑白名單自動(dòng)為算子選擇精度，模型訓(xùn)練過(guò)程中，會(huì)自動(dòng)的插入 cast 操作進(jìn)行類型轉(zhuǎn)換，不需要開(kāi)發(fā)者干預(yù)。針對(duì)數(shù)值精度敏感的計(jì)算，依然使用 FP32 來(lái)算，而對(duì)于數(shù)值安全的計(jì)算，比如矩陣乘，則會(huì)使用  FP16 來(lái)計(jì)算；
• 訓(xùn)練過(guò)程中，會(huì)存儲(chǔ)一份 FP32 的權(quán)重參數(shù)，用于訓(xùn)練時(shí)候的參數(shù)更新，優(yōu)化舍入誤差的問(wèn)題；
• 針對(duì) FP16 容易溢出的問(wèn)題，使用了 Loss scaling 的方案，比如對(duì) Loss 放大 n 倍，根據(jù)鏈?zhǔn)椒▌t，梯度也會(huì)隨之放大 n 倍，從而使其落到 FP16 的表示范圍內(nèi)，具體過(guò)程如下圖左側(cè)所示。

目前所有框架都已經(jīng)支持混合精度。AIAK-Training 組件進(jìn)一步引入 NVIDIA Apex 中  AMP O2 混合精度模式，這個(gè)模式會(huì)更加激進(jìn)的將更多計(jì)算轉(zhuǎn) FP16 來(lái)加速訓(xùn)練。相比默認(rèn) O1 模式，速度會(huì)有進(jìn)一步提升，不過(guò)精度可能會(huì)受影響，需要結(jié)合具體模型驗(yàn)證。

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

下面介紹通信優(yōu)化，這個(gè)也是一個(gè)非常大的話題，涉及到內(nèi)容也非常多。

前面也介紹了，通信主要是分布式訓(xùn)練中引入的，因?yàn)閺膯慰〝U(kuò)展到多卡，需要進(jìn)行多卡之間的一些數(shù)據(jù)同步，這些同步操作就是通過(guò)通信來(lái)實(shí)施的。

下圖列了一個(gè)通信優(yōu)化的整體架構(gòu)：

• 最底層就是網(wǎng)絡(luò)協(xié)議層，包括傳統(tǒng)的 TCP 網(wǎng)絡(luò)、以及在訓(xùn)練場(chǎng)景用的越來(lái)越多的 RoCE 或 IB 的 高性能 RDMA 網(wǎng)絡(luò)。這些底層網(wǎng)絡(luò)方案，百度百舸也都提供了支持。顯然通過(guò)改善硬件基礎(chǔ)設(shè)施來(lái)提升網(wǎng)絡(luò)帶寬、降低延遲，是最直接有效的優(yōu)化通信性能的方法。
• 然后是通信庫(kù)層。因?yàn)樾枰褂玫讓拥木W(wǎng)絡(luò)協(xié)議來(lái)進(jìn)行通信，并且實(shí)際應(yīng)用時(shí)可能涉及多種通信原語(yǔ)，比如點(diǎn)對(duì)點(diǎn)通信、集合通信等，因此也出現(xiàn)一些高度封裝并且優(yōu)化的通信庫(kù)。在 GPU 訓(xùn)練場(chǎng)景中，我們一般都是使用 NCCL 通信庫(kù)，性能比較優(yōu)。
• 基于底層通信庫(kù)，上層框架可以比較方便的構(gòu)建分布式訓(xùn)練的通信架構(gòu)。常見(jiàn)的包括參數(shù)服務(wù)器架構(gòu)，集合通信架構(gòu)。目前 CV/NLP 等領(lǐng)域的模型，主要都是采用集合通信的架構(gòu)方式。
• 通信策略層，主要是在應(yīng)用層上做一些通信效率的優(yōu)化，比如通信隱藏、通信融合、通信壓縮、通信降頻等不同的思路，這些方式大部分也可以疊加一起使用。

我們先看通信策略層面的優(yōu)化思路，首先是通信隱藏優(yōu)化。

在數(shù)據(jù)并行中，梯度同步通信是在訓(xùn)練的反向過(guò)程中進(jìn)行的，當(dāng)反向算出梯度之后，就可以進(jìn)行全局的梯度平均。

如果不做任何的機(jī)制優(yōu)化，反向計(jì)算和通信就會(huì)串行的進(jìn)行，就會(huì)存在計(jì)算上的時(shí)間間隙。

由于反向過(guò)程中，上一個(gè)梯度的通信和下一個(gè)梯度的計(jì)算，兩者之間沒(méi)有任何數(shù)據(jù)依賴，因此可以讓上一個(gè)梯度通信和下一個(gè)梯度計(jì)算并行起來(lái)，讓兩者的耗時(shí)相互重疊，從而可以隱藏部分的通信耗時(shí)。

在實(shí)現(xiàn)層面上，通常是將通信和計(jì)算算子調(diào)度到不同的 cuda 流上實(shí)現(xiàn)的，通信算子調(diào)度到通信流，計(jì)算算子調(diào)度到計(jì)算流，不同流上的算子可以并行發(fā)射執(zhí)行，從而實(shí)現(xiàn)反向中梯度通信和計(jì)算的并行重疊。

目前這個(gè)優(yōu)化能力，在框架中都是默認(rèn)開(kāi)啟的。

其次，通信融合優(yōu)化。

默認(rèn)情況下，模型中的每個(gè)梯度都需要發(fā)起一次通信操作。如果單個(gè)梯度的 size 比較小，那么小數(shù)據(jù)包在實(shí)際通信時(shí)，網(wǎng)絡(luò)帶寬的利用率就會(huì)非常低，通信性能較差。

通信融合，就是將將多個(gè)梯度融合到一起統(tǒng)一進(jìn)行一次通信，從通信開(kāi)銷模型上分析，既能提升帶寬利用率，又能減少通信的初始化延遲項(xiàng)。

現(xiàn)在很多分布式訓(xùn)練框架中，也默認(rèn)都支持梯度融合的策略，不同框架實(shí)現(xiàn)方式有一定的區(qū)別，有的實(shí)現(xiàn)需要先進(jìn)行梯度協(xié)商確定通信順序，有的則是直接通過(guò)靜態(tài)的通信分桶。

雖然框架中默認(rèn)支持通信融合，但是梯度融合的大小一般可以通過(guò)參數(shù)來(lái)配置，用戶可以根據(jù)物理環(huán)境和模型的需求，調(diào)節(jié)合適的融合閾值，應(yīng)該可以取得更佳的收益。

如果在網(wǎng)絡(luò)帶寬較低的訓(xùn)練場(chǎng)景中，比如低帶寬的 TCP 環(huán)境中，梯度同步的延遲可能會(huì)成為訓(xùn)練的主要性能瓶頸。

• 這種情況下，可以考慮的一個(gè)優(yōu)化手段就是通信壓縮。通信壓縮，主要有三種不同的壓縮思路：
• 量化壓縮，比如使用更低精度來(lái)表示梯度，這種方式的壓縮率較低，最大能從  32 位壓縮到 1 位，也就是最大壓縮 32 倍。
• 稀疏化壓縮，典型的算法比如 DGC 算法，核心思想是每輪迭代只傳輸重要的梯度，也就是梯度數(shù)值超過(guò)設(shè)定的某一個(gè)閾值，同時(shí)為了減少信息損失，會(huì)把剩下不重要的梯度在本地進(jìn)行累積，只要時(shí)間足夠，最終累積梯度就會(huì)超過(guò)所設(shè)定的閾值，再進(jìn)行梯度交換。通過(guò)這種方式減少通信的數(shù)據(jù)量，降低對(duì)網(wǎng)絡(luò)帶寬的需求。同時(shí)為了減少對(duì)模型收斂的影響，還通過(guò)動(dòng)量校正、梯度裁剪、動(dòng)量因子掩蔽、預(yù)熱訓(xùn)練等不同方式來(lái)緩解。這種方案目前主要支持 SGD 優(yōu)化器。
• 低秩矩陣壓縮方式，典型的算法比如 PowerSGD，核心思路是將一個(gè)大的梯度矩陣分解成多個(gè)小梯度矩陣，通過(guò)傳輸小矩陣來(lái)減少通信量。

通信降頻優(yōu)化，最簡(jiǎn)單的思路就是增大 batchsize，每次迭代更多的數(shù)據(jù)，減少了迭代次數(shù)，也就是減少了通信量。

不過(guò) batchsize 也不是越大越好，越大的 batchsize 可能會(huì)導(dǎo)致模型收斂精度下降或者收斂速度變慢。針對(duì)類似問(wèn)題，業(yè)界也提出比如 LARS、LAMB 優(yōu)化器算法，通過(guò)分層自適應(yīng)調(diào)節(jié)學(xué)習(xí)率，緩解類似問(wèn)題，AIAK-Training 也增加了支持。

為了增大 batchsize，如果顯存比較充足，可以直接調(diào)整 batchsize 超參。如果顯存比較吃緊，還可以通過(guò)梯度累加的方式，跳過(guò)若干次梯度通信，實(shí)際也相當(dāng)于增大了 batchsize。

下面介紹了一種針對(duì)通信拓?fù)涞膬?yōu)化方案——分層拓?fù)渫ㄐ牛彩轻槍?duì)機(jī)間網(wǎng)絡(luò)帶寬比較低的情況。

通過(guò)分層通信，可以充分的利用機(jī)內(nèi)的高的互聯(lián)帶寬，同時(shí)弱化機(jī)間低網(wǎng)絡(luò)帶寬的影響。

AIAK 中也實(shí)現(xiàn)了這種通信方案，在 25Gbps TCP 環(huán)境下，實(shí)測(cè) 4 機(jī) 32 卡 SwinTransformer 訓(xùn)練，通過(guò)分層 allreduce，性能可以加速 85%。

最后介紹一個(gè)底層通信庫(kù)層面的優(yōu)化，GPU Direct RDMA 通信技術(shù)，這個(gè)技術(shù)需要硬件環(huán)境上支持 RDMA 網(wǎng)絡(luò)。

RDMA 通信，允許本地應(yīng)用程序直接讀寫遠(yuǎn)程應(yīng)用程序的用戶態(tài)虛擬內(nèi)存，整個(gè)通信過(guò)程，除了最一開(kāi)始提交發(fā)送請(qǐng)求這一步需要 CPU 的參與，其余都是由網(wǎng)卡硬件完成的，不需要內(nèi)存拷貝、系統(tǒng)中斷以及軟件處理，因此可以實(shí)現(xiàn)極致的低時(shí)延和高帶寬。

而在 GPU 場(chǎng)景中，GPU Direct RDMA 技術(shù)進(jìn)一步增加了 RDMA 直接訪問(wèn) GPU 顯存的支持，避免通信的時(shí)候，數(shù)據(jù)在 GPU 顯存和主機(jī)內(nèi)存中來(lái)回拷貝，跨機(jī)通信延遲進(jìn)一步降低。

不過(guò)在實(shí)際的案例中，我們發(fā)現(xiàn)有些用戶購(gòu)買了 RDMA 的環(huán)境，但是實(shí)際并沒(méi)有用起來(lái) GDR 技術(shù)，導(dǎo)致通信效率不高。這里列出了幾個(gè)關(guān)鍵的配置項(xiàng)，如果有類似問(wèn)題的話，可以依次進(jìn)行排查和設(shè)置。

前面介紹了當(dāng)前的一些主要的性能優(yōu)化思路和方案，整體來(lái)看，無(wú)論 I/O 優(yōu)化、計(jì)算優(yōu)化、通信優(yōu)化，最樸素的優(yōu)化思路主要就是如何優(yōu)化操作本身、或者是否可以減少操作發(fā)生的次數(shù)，或者是否可以操作和其他的過(guò)程并行起來(lái)從而隱藏開(kāi)銷等。

3. AIAK-Tranining 加速套件實(shí)踐

前面介紹了很多的優(yōu)化工作，要想正確的使能，需要每個(gè)用戶對(duì)于框架的工程實(shí)現(xiàn)原理比較清晰。為了簡(jiǎn)化訓(xùn)練優(yōu)化的成本，我們構(gòu)建了 AIAK-Training 這個(gè)加速套件。

AIAK-Training 會(huì)圍繞數(shù)據(jù)加載、模型計(jì)算、通信等方面，構(gòu)建全鏈路優(yōu)化能力，同時(shí)我們會(huì)把這種優(yōu)化能力，封裝成簡(jiǎn)單易用的接口，用戶插入幾行代碼，即可比較方便的集成使用。同時(shí)，我們也在構(gòu)建自動(dòng)化策略組合調(diào)優(yōu)的機(jī)制，自動(dòng)幫助用戶選擇有效的優(yōu)化策略。

具體使用時(shí)，加速庫(kù)組件可以獨(dú)立的安裝部署，也可以直接使用我們提供的容器鏡像。

下面是一些具體的應(yīng)用案例。

下圖所示，主要是針對(duì) dataloader 的優(yōu)化。這個(gè)場(chǎng)景中，模型比較小，數(shù)據(jù)集規(guī)模也比較小，純計(jì)算的速度其實(shí)比較快，但是跨 EPOCH 的數(shù)據(jù)加載時(shí)間比較長(zhǎng)，導(dǎo)致 I/O 耗時(shí)成為了主要的瓶頸。

通過(guò)使用 AIAK 中提供的進(jìn)程復(fù)用、以及充分預(yù)取機(jī)制，整個(gè)模型訓(xùn)練加速了  166%。

下圖是一個(gè)針對(duì)模型計(jì)算優(yōu)化的案例。

Transformer 類模型訓(xùn)練場(chǎng)景，實(shí)際訓(xùn)練時(shí)通信擴(kuò)展性接近線性，I/O 耗時(shí)占比也非常低，計(jì)算是主要的性能瓶頸。

針對(duì)這個(gè)模型，AIAK-Training 進(jìn)行了一系列計(jì)算層面的優(yōu)化，包括主要結(jié)構(gòu)的算子融合、混合精度、大 batch 調(diào)優(yōu)等等，整個(gè)模型的訓(xùn)練效率提升了 169%。

下圖的案例主要是應(yīng)用了通信層面的優(yōu)化，在云上 TCP 環(huán)境中使能針對(duì)低帶寬網(wǎng)絡(luò)的優(yōu)化策略，在一些經(jīng)典模型上比如 resnet50、bert、vgg16，可以加速 26%~78%。

在自動(dòng)駕駛場(chǎng)景中，我們也針對(duì)典型的 2D 視覺(jué)、3D 視覺(jué)、激光雷達(dá)，以及前融合類的模型，做了一系列的模型訓(xùn)練性能優(yōu)化，訓(xùn)練性能加速 49%~ 391%。