AI 訓(xùn)練加速原理解析與工程實(shí)踐分享
這次分享將系統(tǒng)性的分析在 AI 模型訓(xùn)練過程中的主要性能瓶頸,以及當(dāng)前針對這些瓶頸的主要的加速方案和技術(shù)原理,并介紹百度智能云在這方面的一些實(shí)踐成果。
今天的分享,主要包括三個(gè)部分:
首先介紹我們?yōu)槭裁葱枰?nbsp;AI 訓(xùn)練加速,也就是整體背景和出發(fā)點(diǎn)是什么;
第二部分我們會(huì)系統(tǒng)性的分析實(shí)際訓(xùn)練過程中的可能會(huì)遇到的性能瓶頸問題,然后針對這些問題,介紹目前主要的加速方案;
第三部分介紹百度百舸平臺(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)生模型的過程,用戶需要根據(jù)自己的任務(wù)場景,準(zhǔn)備好訓(xùn)練模型所需要的數(shù)據(jù)集,以及神經(jīng)網(wǎng)絡(luò)算法。
算法可以理解為是一個(gè)高度復(fù)雜的非凸數(shù)學(xué)函數(shù),函數(shù)中包括很多變量以及參數(shù)。模型訓(xùn)練的過程其實(shí)就是在學(xué)習(xí)神經(jīng)網(wǎng)絡(luò)模型中的參數(shù)。
模型訓(xùn)練開始后,會(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í)輸入,通過前向計(jì)算,完成推理。
因此,無論是訓(xùn)練還是推理,核心都是數(shù)據(jù)計(jì)算。 為了加速計(jì)算效率,一般都是通過 GPU 等異構(gòu)加速芯片來進(jìn)行訓(xùn)練和推理。
另外,從深度學(xué)習(xí)模型發(fā)展歷程來看,為了能夠持續(xù)突破模型的精度上限,模型參數(shù)量其實(shí)在快速的膨脹。然而更大的參數(shù)量,就會(huì)帶來更大的計(jì)算復(fù)雜度。
下圖左側(cè)是摘自一篇公開的論文,從這篇總結(jié)里,我們看到在 2010 年之前,模型的計(jì)算量大約 20 個(gè)月翻一番。在 2010~2015 年,常規(guī)模型計(jì)算每 5-6 個(gè)月翻一番。而在 2015 年之后,衍生了大模型訓(xùn)練的趨勢,計(jì)算量增長 10~100 倍。
模型訓(xùn)練對算力以及基礎(chǔ)設(shè)施的要求越來越高,訓(xùn)練需要更多的算力,也需要更長的時(shí)間,這也導(dǎo)致了需要更多的資源成本。這里我們列舉了一些論文或研究中公開的成本數(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 場景提供軟硬一體化的解決方案。通過 AI 計(jì)算、AI 存儲(chǔ)、AI 加速、AI 容器四層技術(shù)棧,滿足上層業(yè)務(wù)場景的需求。
- AI 計(jì)算層,提供了包括高性能的 GPU、以及昆侖等異構(gòu)芯片資源,以及高性能的 RDMA 或 IB 網(wǎng)絡(luò),以及自研的超級 AI 計(jì)算機(jī) X-MAN 等;
- AI 存儲(chǔ)層,包括對象存儲(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ù)公開課有專門分享。
當(dāng)我們考慮做性能加速的時(shí)候,第一個(gè)想到的可能是使用更好的硬件。
這會(huì)帶來一定的性能提升,但是大部分情況下可能并沒有充分發(fā)揮出硬件的計(jì)算能力,核心的原因就是訓(xùn)練代碼的執(zhí)行效率并沒有調(diào)到最優(yōu)或更優(yōu)的狀態(tài)。
- 首先,框架為了兼容更多的需求,一般會(huì)提供一些較為通用的優(yōu)化能力,更多的會(huì)關(guān)注易用性、生態(tài)建設(shè)等;
- 其次,算法工程師在設(shè)計(jì)算法時(shí),核心的精力在如何提高模型的精度,性能效率關(guān)注不足,某些情況下框架提供的一些通用優(yōu)化能力,也都沒有利用起來。
因此,當(dāng)我們決定使用某種模型算法時(shí),為了達(dá)到較好的資源效率和訓(xùn)練效率,我們需要有意識的去優(yōu)化。不過這里也有很多的技術(shù)挑戰(zhàn):
- 性能影響因素比較多,需要根據(jù)不同模型自身的運(yùn)行特點(diǎn)進(jìn)行分析,沒有完全固定的優(yōu)化方案;
- 性能優(yōu)化需要理解工程實(shí)現(xiàn)的原理,比如當(dāng)我們?nèi)プ霎悩?gòu)芯片計(jì)算的優(yōu)化,需要專業(yè)的異構(gòu)研發(fā)經(jīng)驗(yàn)才能開展,技術(shù)門檻較高;
- 還有些情況,當(dāng)我們做分布式訓(xùn)練時(shí),可能還需要結(jié)合集群硬件和網(wǎng)絡(luò)拓?fù)洌瑏韮?yōu)化分布式訓(xùn)練場景下的擴(kuò)展性問題。
這些挑戰(zhàn)極大地影響了模型訓(xùn)練性能的調(diào)優(yōu),因此我們推出了AIAK-Training 的加速套件,期望通過抽象易用性的接口降低優(yōu)化成本,并通過軟硬協(xié)同的優(yōu)化手段,來充分加速客戶在百度智能云上的模型訓(xùn)練性能。
2. 訓(xùn)練性能開銷分析和加速方案
在介紹 AIAK-Training 具體效果之前,我們先介紹下訓(xùn)練加速這個(gè)話題下關(guān)鍵的技術(shù)思路和方案原理是什么樣的。
因?yàn)槟P陀?xùn)練優(yōu)化本身是一個(gè)軟硬件綜合的工作,技術(shù)棧相對比較復(fù)雜,今天的內(nèi)容肯定沒辦法涵蓋全部細(xì)節(jié),我們盡量把關(guān)鍵思路講到。
首先我們看下當(dāng)前的模型訓(xùn)練方案。過去的發(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ù)量相對比較少,單卡的訓(xùn)練時(shí)間可以接受。模型參數(shù)規(guī)模需要保證在訓(xùn)練過程中,單張卡的顯存能夠滿足存儲(chǔ)的上限。按照新的 GPU 卡的顯存容量配置,一般可以放置的最大規(guī)模在10 億級參數(shù)量;
- 當(dāng)數(shù)據(jù)集規(guī)模比較大的時(shí)候,因?yàn)橛?xùn)練過程中需要多次遍歷全量的數(shù)據(jù)集,單卡訓(xùn)練就會(huì)需要比較長的時(shí)間,這時(shí)可以從單卡擴(kuò)展到多卡,通過分布式的方式來加快訓(xùn)練。
這里應(yīng)用最廣泛的就是數(shù)據(jù)并行。在數(shù)據(jù)并行方案下,數(shù)據(jù)集會(huì)被平均切分成多份,然后每張卡上保存完整的模型,各自獨(dú)立并行處理切分后的子數(shù)據(jù)集。
當(dāng)模型參數(shù)量足夠大的時(shí)候,比如參數(shù)量達(dá)到百億、千億級別,單卡放不下完整的模型,這里又出現(xiàn)了模型并行、或者同時(shí)使用數(shù)據(jù)并行和模型并行的混合并行方案。
模型并行,會(huì)將模型切分到不同的卡上,每個(gè)卡上放置模型的一部分,這里又根據(jù)切分的方式不同,比如層內(nèi)切分或?qū)娱g切分,又細(xì)分了 Tensor 并行和流水線并行的方式。
因?yàn)樵谝话隳P陀?xùn)練中,使用更多的還是數(shù)據(jù)并行,我們下面重點(diǎn)還是以數(shù)據(jù)并行為例,來介紹性能優(yōu)化的思路。
我們從軟硬件整體視角先理解下單卡訓(xùn)練過程中存在的性能開銷。
下圖左邊是我們從軟件角度上的訓(xùn)練流程。單卡訓(xùn)練的過程,主要包括數(shù)據(jù)讀取、數(shù)據(jù)預(yù)處理,前向計(jì)算輸出并計(jì)算 loss,根據(jù) loss 函數(shù)反向計(jì)算,得到每一層參數(shù)的梯度,最后根據(jù)梯度更新模型參數(shù)。持續(xù)該過程,直到訓(xùn)練收斂。
下圖右邊,是一個(gè)簡化的節(jié)點(diǎn)硬件拓?fù)鋱D。最上面是數(shù)據(jù)存儲(chǔ),可以是本地存儲(chǔ),也可以是網(wǎng)絡(luò)存儲(chǔ)。然后是 CPU、內(nèi)存,CPU 下通過多個(gè) PCIe Switch 連接著 8 張 GPU 卡,編號從 0~7,8 張卡之間通過 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 開銷。根據(jù)存儲(chǔ)介質(zhì),要讀取的數(shù)據(jù)量不同,這部分時(shí)間開銷也不同;
- 接下來是數(shù)據(jù)預(yù)處理,主要是對讀入的數(shù)據(jù)進(jìn)行一些數(shù)據(jù)增強(qiáng)的操作,比如對圖片進(jìn)行 resize,對比度、飽和度調(diào)整等等,這部分工作大多數(shù)情況下是在 CPU 上。當(dāng)在 CPU 上完成數(shù)據(jù)預(yù)處理操作之后,需要從主機(jī)內(nèi)存拷貝到 GPU 顯存上,涉及到主機(jī)和設(shè)備之間的內(nèi)存拷貝開銷;
- 然后開始前向計(jì)算、反向計(jì)算、參數(shù)更新,這部分更多操作主要是通過 GPU 來進(jìn)行,主要時(shí)間花費(fèi)在 GPU 計(jì)算上面。這個(gè)過程里也可能會(huì)穿插一些 CPU 上的操作,可能也需要做主機(jī)和設(shè)備內(nèi)存之間的拷貝。
因此,從單卡角度看,主要存在 I/O、CPU 預(yù)處理、CPU 和 GPU 之間數(shù)據(jù)拷貝,GPU 計(jì)算等方面的開銷。
接著我們看下數(shù)據(jù)并行的過程。
下圖左邊還是訓(xùn)練的主流程,右邊展示了一個(gè) 3 機(jī) 24 卡的訓(xùn)練集群的硬件拓?fù)洌? 臺(tái)機(jī)器通過網(wǎng)絡(luò)互聯(lián)。
前面的部分我們也介紹到了,在數(shù)據(jù)并行里每個(gè)設(shè)備并行獨(dú)立地執(zhí)行前向和反向計(jì)算過程,因此,每個(gè)訓(xùn)練進(jìn)程也都會(huì)遇到前面講的單卡訓(xùn)練中的性能開銷問題。
數(shù)據(jù)并行為了保證和單卡訓(xùn)練在數(shù)學(xué)上等價(jià),需要確保每張卡的模型參數(shù)在迭代過程中始終保持一致。這里一方面需要讓各 GPU 卡的模型參數(shù)初始化狀態(tài)一致,這個(gè)一般是在訓(xùn)練開始前,通過廣播的方式將第一張卡上的參數(shù)狀態(tài)廣播到其他的卡。
而在訓(xùn)練期間,由于每個(gè)設(shè)備處理的數(shù)據(jù)不同,前向計(jì)算所得到的模型損失值也是不同的,因此還需要在每個(gè)設(shè)備反向計(jì)算出梯度之后,進(jìn)行梯度的平均,以平均后的梯度值來更新模型參數(shù),從而保證每張卡的模型參數(shù)在迭代過程中始終保持一致。
梯度平均涉及到通信的過程,包括節(jié)點(diǎn)內(nèi)部卡之間的通信,以及跨節(jié)點(diǎn)的網(wǎng)絡(luò)通信開銷。這里的通信又包括同步通信和異步通信,不過為了保證模型訓(xùn)練的收斂,一般都是采用同步通信的方案,后續(xù)的優(yōu)化工作也都是基于同步通信的方式來展開。
由上可知,數(shù)據(jù)并行相比單卡訓(xùn)練,主要增加了額外的通信開銷。
通過前述分析,我們知道加速AI 訓(xùn)練不單是某一方面的工作,需要從數(shù)據(jù)加載、模型計(jì)算、分布式通信等系統(tǒng)維度綜合考慮。這里說的數(shù)據(jù)加載,包括數(shù)據(jù) I/O、預(yù)處理、內(nèi)存拷貝等過程。
在具體的優(yōu)化實(shí)踐中,給定一個(gè)待優(yōu)化的模型,加速模型訓(xùn)練就是要不斷提升訓(xùn)練的總體吞吐(每秒可以訓(xùn)練的樣本數(shù))。這個(gè)過程,我們一般可以先分析單卡的訓(xùn)練吞吐,當(dāng)單卡的訓(xùn)練吞吐提升上來之后,我們再看從單卡擴(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)來牽引或衡量我們的工作。
單卡性能優(yōu)化關(guān)鍵包括兩部分:
- 首先數(shù)據(jù)加載效率的優(yōu)化。從存儲(chǔ)系統(tǒng)上,我們可以使用更高性能的存儲(chǔ)介質(zhì),或者基于這些高速存儲(chǔ)介質(zhì)組成的并行文件系統(tǒng),或者說一些緩存加速系統(tǒng)。前面介紹到的,百度百舸也提供了相應(yīng)的存儲(chǔ)系統(tǒng)方案,比如 PFS、RapidFS 等。除此之外,還需要在框架 dataloader 中優(yōu)化數(shù)據(jù)的讀取過程。
- 其次就是模型計(jì)算效率的優(yōu)化。主要考慮如何優(yōu)化計(jì)算實(shí)現(xiàn),怎么提升計(jì)算單元的利用效率,這塊可能就需要結(jié)合模型具體分析。
然后從單卡擴(kuò)展到多卡,目標(biāo)是如何達(dá)到線性加速比。線性加速比這個(gè)指標(biāo),簡單來說就是從 1 張卡擴(kuò)到 2 張卡訓(xùn)練時(shí),訓(xùn)練的性能是否是單卡的2倍。
這里核心就是優(yōu)化分布式的通信效率,一方面是硬件層面的優(yōu)化,另外一方面在實(shí)際通信中,需要考慮怎么利用好網(wǎng)絡(luò)的帶寬資源,或者是否能夠?qū)⑼ㄐ胚^程進(jìn)行隱藏等。
下面我們分別從這幾個(gè)方面詳細(xì)展開。
首先是數(shù)據(jù)加載方面的優(yōu)化。
當(dāng)我們實(shí)例化一個(gè) dataloader 之后,我們會(huì)持續(xù)迭代 dataloader 來讀取一個(gè) batch 的數(shù)據(jù)進(jìn)行模型訓(xùn)練。
如果這里不做任何優(yōu)化,如下圖上半部分所示,每個(gè) batch 的數(shù)據(jù)加載過程和每個(gè) batch 的模型訓(xùn)練過程,實(shí)際上是串行進(jìn)行的。從 GPU 視角來看,就會(huì)出現(xiàn)因?yàn)閿?shù)據(jù)加載導(dǎo)致的計(jì)算間隙,計(jì)算資源存在時(shí)間上的浪費(fèi)。
除了前面講到我們用更好的硬件直接提升數(shù)據(jù)讀的效率之外,還可以怎么優(yōu)化呢?
實(shí)際在 AI 訓(xùn)練過程中,數(shù)據(jù)訪問上有兩個(gè)關(guān)鍵特征:
- 當(dāng)數(shù)據(jù)集做完 shuffle 之后,每輪訓(xùn)練所需要的 batch 數(shù)據(jù)以及訪問順序是已知的;
- 任意兩個(gè) batch 的數(shù)據(jù)讀可以并行,因?yàn)閿?shù)據(jù)之間沒有任何依賴關(guān)系。
因此,我們在不依賴硬件層面改動(dòng)的時(shí)候,可以做的一個(gè)優(yōu)化工作就是數(shù)據(jù)預(yù)取,當(dāng)訓(xùn)練第一個(gè) batch 數(shù)據(jù)的時(shí)候,可以提前加載下一個(gè) batch 數(shù)據(jù),讓 I/O 的過程和 GPU 上的計(jì)算充分并行起來。
首先,我們需要利用好 dataloader 中已有的優(yōu)化方案,一是合理設(shè)置 num_workers 超參數(shù),通過多進(jìn)程的方式讀數(shù)據(jù),這一步可以實(shí)現(xiàn)數(shù)據(jù)從存儲(chǔ)系統(tǒng)中預(yù)取到主機(jī)內(nèi)存。二是從主機(jī)內(nèi)存拷貝到 GPU顯存,可以通過 pinned memory 的機(jī)制來加速。
大概介紹下 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)的操作。
不過當(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 針對這個(gè)問題,又做了進(jìn)一步的優(yōu)化,可以實(shí)現(xiàn) H2D 和前向計(jì)算的 overlap。
在數(shù)據(jù)并行場景下,還需要注意的一個(gè)事情,數(shù)據(jù)需要均衡的切分。
如果每個(gè)訓(xùn)練進(jìn)程分配的數(shù)據(jù)不均衡,計(jì)算量就會(huì)不同,也就導(dǎo)致每個(gè)進(jìn)程前向計(jì)算和反向計(jì)算完成的時(shí)間不同,那么先完成計(jì)算的進(jìn)程,在反向過程中就會(huì)先進(jìn)入到梯度通信環(huán)節(jié)中,但是因?yàn)?Allreduce 通信是同步通信操作,需要所有進(jìn)程同時(shí)開始并同時(shí)結(jié)束,因此先開始通信的進(jìn)程,會(huì)一直等待其他所有進(jìn)程也發(fā)起了 AllReduce 后才能一起完成通信操作。這里就會(huì)出現(xiàn)因?yàn)榭炻灰粚?dǎo)致的資源空閑問題。
為了解決這種問題,需要每個(gè)進(jìn)程使用相同的 batchsize 來讀取數(shù)據(jù),另外每個(gè) batch 的數(shù)據(jù)量要均衡。圖像類數(shù)據(jù)一般會(huì)固定尺寸進(jìn)行訓(xùn)練,而像 NLP 類模型需要處理變長的語句,可能需要進(jìn)行特殊的處理,比如可以將數(shù)據(jù) padding 到相同的長度,或者通過樣本長度排序的方式來均衡分配等。
下面介紹計(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í)候,簡化的流程有四步。
- 首先在 CPU 上異步發(fā)射一個(gè) GPU 計(jì)算 Kernel;
- 當(dāng) Kernel 調(diào)度執(zhí)行時(shí),需要先從 GPU 上的 Global Memory 讀取計(jì)算所需要的數(shù)據(jù);
- 開始執(zhí)行計(jì)算;
- 當(dāng)計(jì)算完成后,需要將計(jì)算的結(jié)果寫回 Global Memory。
根據(jù)計(jì)算、訪存的開銷占比,一般會(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ì)算間隙問題。
由上可知,優(yōu)化模型計(jì)算效率,需要從訪存優(yōu)化、計(jì)算優(yōu)化、其他開銷的優(yōu)化綜合考慮。
- 訪存優(yōu)化,主要考慮如何減少數(shù)據(jù)在顯存和計(jì)算單元之間搬運(yùn)的時(shí)間。從算子實(shí)現(xiàn)角度上,需要利用好 GPU 存儲(chǔ)層次架構(gòu),比如把數(shù)據(jù)搬運(yùn)到更快的存儲(chǔ)器上比如 share memory,減少對 global memory 的訪問,從而節(jié)省訪存時(shí)間。或者做好計(jì)算指令和訪存指令的 overlap,來提升計(jì)算訪存比。而從單算子擴(kuò)展到多算子時(shí),還需要考慮如何減少中間結(jié)果的讀寫,這種問題一般可以通過算子融合的手段來優(yōu)化;
- 計(jì)算優(yōu)化,計(jì)算瓶頸問題多半應(yīng)該是沒有正確對任務(wù)進(jìn)行分塊,或者沒有利用好 GPU 并行計(jì)算的優(yōu)勢,導(dǎo)致并行度不高。還有可能沒有使用合并指令集導(dǎo)致計(jì)算效率低下,或者沒有使用 Tensor Core 等高性能計(jì)算單元,造成資源浪費(fèi)。不同的問題,也對應(yīng)著不同的優(yōu)化手段;
- Kernel Launch 等其他開銷,大量時(shí)間花費(fèi)在訪存或計(jì)算之外,可以考慮的優(yōu)化手段如算子融合、Cuda Graph 等。
首先是算子融合。算子在底層 GPU 執(zhí)行時(shí),會(huì)發(fā)起一次或者多次 Kernel Launch,Kernel 之間交互數(shù)據(jù)也需要經(jīng)過顯存,而算子融合就是將多個(gè) GPU Kernel 融合成一個(gè)大 Kernel,統(tǒng)一發(fā)起和執(zhí)行。
- 因?yàn)闇p少了需要執(zhí)行的算子數(shù)量,從而可以減少 Kernel 調(diào)度和發(fā)起的開銷;
- 通過融合,可以通過寄存器等來傳遞中間結(jié)果,避免從 global memory 的來回搬運(yùn),極大降低了顯存等待的時(shí)間;
- 在某些場景中,可以通過算子融合,可以更充分的利用計(jì)算資源,提升資源和計(jì)算的效率。
算子融合具體如何實(shí)現(xiàn)呢?
一種方式,分析模型中的低效操作,專家經(jīng)驗(yàn)手寫融合算子。在 GPU 上主要就是 CUDA 算子研發(fā),這里存在一定的門檻。AIAK-Training 會(huì)針對典型的模型結(jié)構(gòu),或者客戶需求,提供高效優(yōu)化的算子實(shí)現(xiàn)。
另一種方式,就是編譯優(yōu)化的方案。通過編譯的方式進(jìn)行計(jì)算優(yōu)化,以及代碼自動(dòng)生成,從而降低在不同硬件上的手工優(yōu)化成本。不過當(dāng)前很多編譯方案更多還是針對推理優(yōu)化,訓(xùn)練上的方案還在快速演進(jìn)過程中。不過從極致性能角度看,未來一段時(shí)間依然離不開手寫融合算子的工作。
下面介紹一些算子融合的實(shí)際案例。第一個(gè)是針對典型模型網(wǎng)絡(luò)結(jié)構(gòu)的優(yōu)化。
下圖展示了我們在 SwinTransformer 模型中針對核心模塊 WindowAttention 進(jìn)行的計(jì)算融合優(yōu)化。
WindowAttention 結(jié)構(gòu),核心操作公式如下圖所示。計(jì)算過程中,需要依次執(zhí)行 7 個(gè)計(jì)算 Kernel。再加上一些 reshape 等轉(zhuǎn)換操作,總共需要 launch 10 個(gè) Kernel。通過性能分析發(fā)現(xiàn),實(shí)際執(zhí)行過程中 launch kernel 的間隔冗余開銷占到了端到端 80% 以上的時(shí)間,導(dǎo)致該模塊存在著較大的優(yōu)化空間。
通過將這些 Kernel 融合成一個(gè),整個(gè)模塊的執(zhí)行時(shí)間從 392 微秒減少到 13 微秒,單算子加速了 30 倍。整個(gè)模型的訓(xùn)練效率,端到端加速了 20% 以上。
這個(gè)優(yōu)化的核心思路,主要有三點(diǎn):
- 是利用好 GPU 的三級訪存流水:顯存、share memory、寄存器;
- 通過分塊策略,將 2 個(gè)矩陣乘法和 softmax 進(jìn)行融合;
- 針對前向中間結(jié)果的優(yōu)化,充分利用 Tensor Core,在反向計(jì)算中使用重計(jì)算代替重加載,使得訪存開銷極大降低。
下圖是一個(gè)數(shù)據(jù)操作的融合舉例,是在 FCOS3D 模型中對于坐標(biāo)壓縮操作的一個(gè)優(yōu)化。
通過性能分析發(fā)現(xiàn),這個(gè)操作過程中存在大量 GPU 空隙,GPU 利用率較低。該操作主要的功能是根據(jù) index 將 3D-Tensor 壓縮成 2D-Tensor。在原生實(shí)現(xiàn)中會(huì)先在host 端生成 index,再進(jìn)行 H2D 拷貝,最后完成 Tensor 壓縮,這會(huì)造成額外的拷貝和等待開銷。
為此,我們重新實(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í)際案例來介紹。
我們發(fā)現(xiàn)在一些模型中,有一些操作是串行執(zhí)行的。比如在一些目標(biāo)檢測模型里,在 loss 計(jì)算過程中,有些操作并不是按照一個(gè) batch 進(jìn)行操作,而是 for-loop 每張圖片或一個(gè)樣本,這種情況下,當(dāng)我們?nèi)ヌ嵘?nbsp;batchsize 的時(shí)候,因?yàn)檫@里的串行,可能沒法達(dá)到我們想要的性能效果。
以 YOLOv7 中的 SimOTA 操作舉例,原生實(shí)現(xiàn)中,是通過 for-loop 遍歷一個(gè) batch 的每一張圖片,然后為圖片的 gtbox 執(zhí)行 SimOTA 標(biāo)簽分配。這種串行的實(shí)現(xiàn)方式導(dǎo)致該部分操作的 GPU 利用率非常低效。
而每張圖片處理時(shí)數(shù)據(jù)之間是沒有依賴的。因此,我們做的一項(xiàng)工作就是將串行計(jì)算改成 batch 并行計(jì)算,通過對一個(gè) batch 的數(shù)據(jù)進(jìn)行并行化的標(biāo)簽分配,從而加速這部分計(jì)算的效率。
最終效果上,SimOTA 操作的耗時(shí)從 384 毫秒下降到 69 毫秒,計(jì)算效率提升 5.5 倍,整個(gè)模型的端到端訓(xùn)練效率提升了 18% 以上。
模型訓(xùn)練過程中也有其他的類似場景,比如說參數(shù)更新。參數(shù)更新時(shí),默認(rèn)也是通過循環(huán)的方式,遍歷每個(gè)參數(shù),然后每個(gè)參數(shù)都會(huì)啟動(dòng)一個(gè)參數(shù)更新的 Cuda Kernel,然后依次執(zhí)行。
針對這種情況,AIAK 也增加了 FusedOptimizer 的優(yōu)化,通過融合參數(shù)更新的算子,實(shí)現(xiàn)批量化更新參數(shù),大幅減少 Kernel Launch 次數(shù)。
下面介紹了另一個(gè)優(yōu)化手段 CUDA Graph,主要是為了減少 CPU Launch Kernel 的開銷。
CUDA Graph 是在 CUDA 10 版本中引入的特性,可以將一系列 CUDA Kernel 封裝成單個(gè)單元,可以通過一次 CPU Launch 操作來啟動(dòng)多個(gè) GPU Kernel,從而減少了CPU Launch Kernel 的開銷。
如下圖所示,默認(rèn)情況下 CPU 需要依次發(fā)射多個(gè) Kernel,如果 Kernel 計(jì)算時(shí)間比較短,Kernel 之間的 Launch 間隙可能成為性能瓶頸。通過 CUDA Graph,僅需要花費(fèi)一些額外時(shí)間構(gòu)建 Graph,后續(xù)通過一次 Graph 的發(fā)射,即可大幅縮短實(shí)際執(zhí)行時(shí) Kernel 之間的間隙。
現(xiàn)在很多框架也都增加了對 CUDA Graph 的支持,通過插入一些代碼來使K能這個(gè)功能。不過也有一些使用上的限制,比如不支持動(dòng)態(tài) shape、不支持動(dòng)態(tài)控制流、過程中不能捕獲 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 開始引入,主要是用來加速矩陣或卷積的操作效率。
相比 FP32 CUDA Core 一次只能執(zhí)行兩個(gè)標(biāo)量的計(jì)算,Tensor Core 可以一次對兩個(gè)矩陣執(zhí)行計(jì)算,因此 Tensor Core 的計(jì)算吞吐遠(yuǎn)高于 FP32 CUDA Core。
在 A100 中,Tensor Core 支持了多種浮點(diǎn)數(shù)據(jù)類型,對于深度學(xué)習(xí)訓(xùn)練來說,可能涉及到包括 FP16、BF16、以及 TF32 模式。
TF32 主要用于單精度訓(xùn)練場景,相比 FP32 訓(xùn)練,在保持相同的訪存帶寬需求下,理論計(jì)算吞吐量提高了 8 倍。
FP16/BF16 主要用于混合精度訓(xùn)練場景,相比 FP32 訓(xùn)練,訪存需求減少了一半,理論計(jì)算吞吐量提高了 16 倍。
使用 Tensor Core 的方式,可以使用底層的 cublas 或 cuda 接口進(jìn)行編程,而對于算法開發(fā)者來說,更直接的就是使用框架中提供的 TF32 訓(xùn)練或混合精度訓(xùn)練方案。
首先是 TF32 訓(xùn)練模式,TF32 是 Ampere 開始引入。
TF32 在浮點(diǎn)數(shù)的表達(dá)中,有 8 個(gè)指數(shù)位,10 個(gè)尾數(shù)位和 1 個(gè)符號位。指數(shù)位與 FP32 相同,即數(shù)據(jù)表示范圍相同,但是尾數(shù)位低于 FP32,和 FP16 相同。
需要注意的是,TF32 不是一個(gè)對外開放的數(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)開啟,某些框架版本中可能需要通過環(huán)境變量或者參數(shù)配置來手工開啟,具體需要參考框架的用戶手冊。
不過由于 TF32 相比 FP32 來說,精度范圍降低了,實(shí)際訓(xùn)練時(shí)還需要關(guān)注對模型收斂精度的影響。
混合精度訓(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)算效率更快;
不過因?yàn)?FP16 的指數(shù)位和尾數(shù)位的范圍都比 FP32 要少,因此數(shù)值表示范圍和精度都會(huì)有降低,在實(shí)際使用的時(shí)候,就可能出現(xiàn)因?yàn)楸硎痉秶M窄導(dǎo)致的數(shù)值溢出問題,或者因?yàn)榫炔蛔銓?dǎo)致舍入誤差。
為了優(yōu)化類似問題,混合精度訓(xùn)練方案中有幾個(gè)關(guān)鍵的技術(shù)工作:
- 算子黑白名單機(jī)制,框架使用黑白名單自動(dòng)為算子選擇精度,模型訓(xùn)練過程中,會(huì)自動(dòng)的插入 cast 操作進(jìn)行類型轉(zhuǎn)換,不需要開發(fā)者干預(yù)。針對數(shù)值精度敏感的計(jì)算,依然使用 FP32 來算,而對于數(shù)值安全的計(jì)算,比如矩陣乘,則會(huì)使用 FP16 來計(jì)算;
- 訓(xùn)練過程中,會(huì)存儲(chǔ)一份 FP32 的權(quán)重參數(shù),用于訓(xùn)練時(shí)候的參數(shù)更新,優(yōu)化舍入誤差的問題;
- 針對 FP16 容易溢出的問題,使用了 Loss scaling 的方案,比如對 Loss 放大 n 倍,根據(jù)鏈?zhǔn)椒▌t,梯度也會(huì)隨之放大 n 倍,從而使其落到 FP16 的表示范圍內(nèi),具體過程如下圖左側(cè)所示。
目前所有框架都已經(jīng)支持混合精度。AIAK-Training 組件進(jìn)一步引入 NVIDIA Apex 中 AMP O2 混合精度模式,這個(gè)模式會(huì)更加激進(jìn)的將更多計(jì)算轉(zhuǎn) FP16 來加速訓(xùn)練。相比默認(rèn) O1 模式,速度會(huì)有進(jìn)一步提升,不過精度可能會(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ù)同步,這些同步操作就是通過通信來實(shí)施的。
下圖列了一個(gè)通信優(yōu)化的整體架構(gòu):
- 最底層就是網(wǎng)絡(luò)協(xié)議層,包括傳統(tǒng)的 TCP 網(wǎng)絡(luò)、以及在訓(xùn)練場景用的越來越多的 RoCE 或 IB 的 高性能 RDMA 網(wǎng)絡(luò)。這些底層網(wǎng)絡(luò)方案,百度百舸也都提供了支持。顯然通過改善硬件基礎(chǔ)設(shè)施來提升網(wǎng)絡(luò)帶寬、降低延遲,是最直接有效的優(yōu)化通信性能的方法。
- 然后是通信庫層。因?yàn)樾枰褂玫讓拥木W(wǎng)絡(luò)協(xié)議來進(jìn)行通信,并且實(shí)際應(yīng)用時(shí)可能涉及多種通信原語,比如點(diǎn)對點(diǎn)通信、集合通信等,因此也出現(xiàn)一些高度封裝并且優(yōu)化的通信庫。在 GPU 訓(xùn)練場景中,我們一般都是使用 NCCL 通信庫,性能比較優(yōu)。
- 基于底層通信庫,上層框架可以比較方便的構(gòu)建分布式訓(xùn)練的通信架構(gòu)。常見的包括參數(shù)服務(wù)器架構(gòu),集合通信架構(gòu)。目前 CV/NLP 等領(lǐng)域的模型,主要都是采用集合通信的架構(gòu)方式。
- 通信策略層,主要是在應(yīng)用層上做一些通信效率的優(yōu)化,比如通信隱藏、通信融合、通信壓縮、通信降頻等不同的思路,這些方式大部分也可以疊加一起使用。
我們先看通信策略層面的優(yōu)化思路,首先是通信隱藏優(yōu)化。
在數(shù)據(jù)并行中,梯度同步通信是在訓(xùn)練的反向過程中進(jìn)行的,當(dāng)反向算出梯度之后,就可以進(jìn)行全局的梯度平均。
如果不做任何的機(jī)制優(yōu)化,反向計(jì)算和通信就會(huì)串行的進(jìn)行,就會(huì)存在計(jì)算上的時(shí)間間隙。
由于反向過程中,上一個(gè)梯度的通信和下一個(gè)梯度的計(jì)算,兩者之間沒有任何數(shù)據(jù)依賴,因此可以讓上一個(gè)梯度通信和下一個(gè)梯度計(jì)算并行起來,讓兩者的耗時(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)開啟的。
其次,通信融合優(yōu)化。
默認(rèn)情況下,模型中的每個(gè)梯度都需要發(fā)起一次通信操作。如果單個(gè)梯度的 size 比較小,那么小數(shù)據(jù)包在實(shí)際通信時(shí),網(wǎng)絡(luò)帶寬的利用率就會(huì)非常低,通信性能較差。
通信融合,就是將將多個(gè)梯度融合到一起統(tǒng)一進(jìn)行一次通信,從通信開銷模型上分析,既能提升帶寬利用率,又能減少通信的初始化延遲項(xiàng)。
現(xiàn)在很多分布式訓(xùn)練框架中,也默認(rèn)都支持梯度融合的策略,不同框架實(shí)現(xiàn)方式有一定的區(qū)別,有的實(shí)現(xiàn)需要先進(jìn)行梯度協(xié)商確定通信順序,有的則是直接通過靜態(tài)的通信分桶。
雖然框架中默認(rèn)支持通信融合,但是梯度融合的大小一般可以通過參數(shù)來配置,用戶可以根據(jù)物理環(huán)境和模型的需求,調(diào)節(jié)合適的融合閾值,應(yīng)該可以取得更佳的收益。
如果在網(wǎng)絡(luò)帶寬較低的訓(xùn)練場景中,比如低帶寬的 TCP 環(huán)境中,梯度同步的延遲可能會(huì)成為訓(xùn)練的主要性能瓶頸。
- 這種情況下,可以考慮的一個(gè)優(yōu)化手段就是通信壓縮。通信壓縮,主要有三種不同的壓縮思路:
- 量化壓縮,比如使用更低精度來表示梯度,這種方式的壓縮率較低,最大能從 32 位壓縮到 1 位,也就是最大壓縮 32 倍。
- 稀疏化壓縮,典型的算法比如 DGC 算法,核心思想是每輪迭代只傳輸重要的梯度,也就是梯度數(shù)值超過設(shè)定的某一個(gè)閾值,同時(shí)為了減少信息損失,會(huì)把剩下不重要的梯度在本地進(jìn)行累積,只要時(shí)間足夠,最終累積梯度就會(huì)超過所設(shè)定的閾值,再進(jìn)行梯度交換。通過這種方式減少通信的數(shù)據(jù)量,降低對網(wǎng)絡(luò)帶寬的需求。同時(shí)為了減少對模型收斂的影響,還通過動(dòng)量校正、梯度裁剪、動(dòng)量因子掩蔽、預(yù)熱訓(xùn)練等不同方式來緩解。這種方案目前主要支持 SGD 優(yōu)化器。
- 低秩矩陣壓縮方式,典型的算法比如 PowerSGD,核心思路是將一個(gè)大的梯度矩陣分解成多個(gè)小梯度矩陣,通過傳輸小矩陣來減少通信量。
通信降頻優(yōu)化,最簡單的思路就是增大 batchsize,每次迭代更多的數(shù)據(jù),減少了迭代次數(shù),也就是減少了通信量。
不過 batchsize 也不是越大越好,越大的 batchsize 可能會(huì)導(dǎo)致模型收斂精度下降或者收斂速度變慢。針對類似問題,業(yè)界也提出比如 LARS、LAMB 優(yōu)化器算法,通過分層自適應(yīng)調(diào)節(jié)學(xué)習(xí)率,緩解類似問題,AIAK-Training 也增加了支持。
為了增大 batchsize,如果顯存比較充足,可以直接調(diào)整 batchsize 超參。如果顯存比較吃緊,還可以通過梯度累加的方式,跳過若干次梯度通信,實(shí)際也相當(dāng)于增大了 batchsize。
下面介紹了一種針對通信拓?fù)涞膬?yōu)化方案——分層拓?fù)渫ㄐ?,也是針對機(jī)間網(wǎng)絡(luò)帶寬比較低的情況。
通過分層通信,可以充分的利用機(jī)內(nèi)的高的互聯(lián)帶寬,同時(shí)弱化機(jī)間低網(wǎng)絡(luò)帶寬的影響。
AIAK 中也實(shí)現(xiàn)了這種通信方案,在 25Gbps TCP 環(huán)境下,實(shí)測 4 機(jī) 32 卡 SwinTransformer 訓(xùn)練,通過分層 allreduce,性能可以加速 85%。
最后介紹一個(gè)底層通信庫層面的優(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è)通信過程,除了最一開始提交發(fā)送請求這一步需要 CPU 的參與,其余都是由網(wǎng)卡硬件完成的,不需要內(nèi)存拷貝、系統(tǒng)中斷以及軟件處理,因此可以實(shí)現(xiàn)極致的低時(shí)延和高帶寬。
而在 GPU 場景中,GPU Direct RDMA 技術(shù)進(jìn)一步增加了 RDMA 直接訪問 GPU 顯存的支持,避免通信的時(shí)候,數(shù)據(jù)在 GPU 顯存和主機(jī)內(nèi)存中來回拷貝,跨機(jī)通信延遲進(jìn)一步降低。
不過在實(shí)際的案例中,我們發(fā)現(xiàn)有些用戶購買了 RDMA 的環(huán)境,但是實(shí)際并沒有用起來 GDR 技術(shù),導(dǎo)致通信效率不高。這里列出了幾個(gè)關(guān)鍵的配置項(xiàng),如果有類似問題的話,可以依次進(jìn)行排查和設(shè)置。
前面介紹了當(dāng)前的一些主要的性能優(yōu)化思路和方案,整體來看,無論 I/O 優(yōu)化、計(jì)算優(yōu)化、通信優(yōu)化,最樸素的優(yōu)化思路主要就是如何優(yōu)化操作本身、或者是否可以減少操作發(fā)生的次數(shù),或者是否可以操作和其他的過程并行起來從而隱藏開銷等。
3. AIAK-Tranining 加速套件實(shí)踐
前面介紹了很多的優(yōu)化工作,要想正確的使能,需要每個(gè)用戶對于框架的工程實(shí)現(xiàn)原理比較清晰。為了簡化訓(xùn)練優(yōu)化的成本,我們構(gòu)建了 AIAK-Training 這個(gè)加速套件。
AIAK-Training 會(huì)圍繞數(shù)據(jù)加載、模型計(jì)算、通信等方面,構(gòu)建全鏈路優(yōu)化能力,同時(shí)我們會(huì)把這種優(yōu)化能力,封裝成簡單易用的接口,用戶插入幾行代碼,即可比較方便的集成使用。同時(shí),我們也在構(gòu)建自動(dòng)化策略組合調(diào)優(yōu)的機(jī)制,自動(dòng)幫助用戶選擇有效的優(yōu)化策略。
具體使用時(shí),加速庫組件可以獨(dú)立的安裝部署,也可以直接使用我們提供的容器鏡像。
下面是一些具體的應(yīng)用案例。
下圖所示,主要是針對 dataloader 的優(yōu)化。這個(gè)場景中,模型比較小,數(shù)據(jù)集規(guī)模也比較小,純計(jì)算的速度其實(shí)比較快,但是跨 EPOCH 的數(shù)據(jù)加載時(shí)間比較長,導(dǎo)致 I/O 耗時(shí)成為了主要的瓶頸。
通過使用 AIAK 中提供的進(jìn)程復(fù)用、以及充分預(yù)取機(jī)制,整個(gè)模型訓(xùn)練加速了 166%。
下圖是一個(gè)針對模型計(jì)算優(yōu)化的案例。
Transformer 類模型訓(xùn)練場景,實(shí)際訓(xùn)練時(shí)通信擴(kuò)展性接近線性,I/O 耗時(shí)占比也非常低,計(jì)算是主要的性能瓶頸。
針對這個(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)境中使能針對低帶寬網(wǎng)絡(luò)的優(yōu)化策略,在一些經(jīng)典模型上比如 resnet50、bert、vgg16,可以加速 26%~78%。
在自動(dòng)駕駛場景中,我們也針對典型的 2D 視覺、3D 視覺、激光雷達(dá),以及前融合類的模型,做了一系列的模型訓(xùn)練性能優(yōu)化,訓(xùn)練性能加速 49%~ 391%。