從「根」上找出模型的瓶頸！從第一原理出發(fā)剖析深度學(xué)習(xí)

作者：新智元 2022-03-24 14:51:33

了解事物的底層邏輯才能更好地解決問題。康奈爾大學(xué)AI聯(lián)合創(chuàng)始人最近發(fā)了一篇文章，從第一原理出發(fā)，深度剖析深度學(xué)習(xí)性能瓶頸的三座大山：計(jì)算、內(nèi)存和開銷。

如果想提升模型的性能，你的第一直覺是問搜索引擎嗎？

通常情況下你得到的建議只能是一些技巧性的操作，比如使用in-place operation，把梯度設(shè)置為None，或者是把PyTorch版本從1.10.1退回到穩(wěn)定版1.10.0等等。

這些臨時(shí)找到的騷操作雖然可以一時(shí)地解決當(dāng)下問題，但要是用了以后性能還沒提升到滿意的程度，那可能就有點(diǎn)「抓瞎」了。

雖然深度學(xué)習(xí)本身就是一個積木類的黑盒模型，但這種調(diào)試方法仿佛深度學(xué)習(xí)真的變成了煉丹術(shù)，而非科學(xué)。

比如你的模型在訓(xùn)練集上的loss遠(yuǎn)遠(yuǎn)低于測試時(shí)的loss，說明模型已經(jīng)「過擬合」了，如果這個時(shí)候再盲目增大模型的參數(shù)量，那就純粹是浪費(fèi)時(shí)間了。再比如模型的訓(xùn)練loss和驗(yàn)證loss一樣的時(shí)候，如果再對模型加入正則化，那也是浪費(fèi)時(shí)間。

所以為了讓AI從業(yè)者在遇到問題之后，能從根上解決，最近康奈爾大學(xué)人工智能（CUAI）的一位聯(lián)合創(chuàng)始人Horace He發(fā)表了一篇博客，把深度學(xué)習(xí)模型的時(shí)間損耗拆分成三部分：計(jì)算、內(nèi)存和其他開銷overhead，從「第一原理」出發(fā)來了解和改進(jìn)深度學(xué)習(xí)模型。

其中計(jì)算（Compute）指的是GPU在計(jì)算浮點(diǎn)操作時(shí)所消耗的時(shí)間，也就是FLOPS；內(nèi)存（Memory）指的是把tensors寫到GPU里消耗的時(shí)間。

如果模型把大部分的時(shí)間都花在了內(nèi)存?zhèn)鬏斏?，那么增加GPU的FLOPS是沒有用的。又或者如果你把所有的時(shí)間都花在執(zhí)行大塊的數(shù)學(xué)運(yùn)算上，那么把你的模型邏輯改寫成C++來減少開銷也沒有用。

了解你所處的狀態(tài)可以讓你縮小優(yōu)化的范圍，節(jié)省下來的時(shí)間就可以愉快地摸魚了。

計(jì)算

通常深度學(xué)習(xí)模型運(yùn)算速度不夠快的原因都是顯卡性能不夠，加卡解千愁?。?/span>

但現(xiàn)實(shí)很骨感，越強(qiáng)的卡，價(jià)格也更美麗。所以為了錢花的更值，需要盡可能地提升顯卡的運(yùn)行效率，不斷地讓顯卡進(jìn)行矩陣運(yùn)行。

并且計(jì)算比內(nèi)存帶寬更重要的原因還有一個，就是模型訓(xùn)練過程中所需的計(jì)算量不管通過何種手段，基本都不會降低，所以最大限度提升計(jì)算能力才能提升效率。

但計(jì)算量如果增長速度過快，也會加劇最大化計(jì)算利用率的難度。就拿這個關(guān)于CPU FLOPS翻倍時(shí)間與內(nèi)存帶寬翻倍時(shí)間的表格來說。

一種思考計(jì)算的方式是把CPU當(dāng)作一個工廠。用戶向工廠發(fā)送指令（開銷）和原材料（內(nèi)存帶寬），所有這些都是為了保持工廠高效運(yùn)行（計(jì)算）。

如果工廠提高效率的速度超過了為其提供原材料的速度，那么工廠就更難達(dá)到其峰值效率。即使工廠的規(guī)模（FLOPS）增加了一倍，如果帶寬不能同步提升，那性能也不會增加一倍。

關(guān)于FLOPS還有一個補(bǔ)充?，F(xiàn)代機(jī)器學(xué)習(xí)加速硬件都有專門用于矩陣乘法的硬件，比如Nvidia的Tensor Cores。

也就是說，如果你不做矩陣乘法，你就只能獲得19.5 teraflops，而非宣傳的312。并且這并非是GPU所獨(dú)有的缺陷，TPU甚至比GPU更不通用。

事實(shí)上，GPU在所有非矩陣乘法的操作上都很慢，乍一看可能影響很大，但實(shí)際上神經(jīng)網(wǎng)絡(luò)模型里基本都是矩陣乘法。

在一篇關(guān)于BERT模型的flop研究中可以發(fā)現(xiàn)，BERT中99.8%都是矩陣乘法（Tensor Contraction）操作，所以雖然非矩陣乘法的速度要慢15倍，但也無傷大雅。

但在這種情況下，歸一化和點(diǎn)式運(yùn)算實(shí)際上比矩陣乘法運(yùn)算少了250倍的FLOPS和700倍的FLOPS。

至于為什么非矩陣乘法的理論性能和現(xiàn)實(shí)相差這么多，研究人員給出的答案是：內(nèi)存帶寬（memory bandwidth）。

內(nèi)存

帶寬成本本質(zhì)上是將數(shù)據(jù)從一個地方移動到另一個地方所支付的成本，包括將數(shù)據(jù)從CPU轉(zhuǎn)移到GPU，從一個節(jié)點(diǎn)轉(zhuǎn)移到另一個節(jié)點(diǎn)，二者通常稱為「數(shù)據(jù)傳輸成本」和「網(wǎng)絡(luò)成本」。

深度學(xué)習(xí)模型優(yōu)化關(guān)注的帶寬成本主要是從CUDA全局內(nèi)存轉(zhuǎn)移到CUDA共享內(nèi)存。

回到工廠那個例子，雖然工廠可以完成一些計(jì)算任務(wù)，但它并不是一個適合存儲大量數(shù)據(jù)的地方。典型的做法是利用更便宜的硬件來建立一個數(shù)據(jù)倉庫（DRAM），然后在倉庫和工廠之間運(yùn)送物資，也就是內(nèi)存帶寬。

GPU的DRAM大小可以通過nvidia-smi命令獲得，倉庫容量不夠也是導(dǎo)致CUDA Out of Memory錯誤的主要原因。

需要注意的是，每次執(zhí)行GPU內(nèi)核時(shí)，都需要將數(shù)據(jù)從GPU的DRAM移出和移回。

現(xiàn)在我們就知道執(zhí)行torch.cos這樣的單個操作時(shí)，幾乎每做一次這樣的簡單運(yùn)算，數(shù)據(jù)都需要從內(nèi)存運(yùn)到GPU里，運(yùn)送成本比計(jì)算成本要高很多，所以時(shí)間幾乎都花在內(nèi)存上了，這種情況也稱為memory-bound operation。

錯誤的做法就是每次都把數(shù)據(jù)送到GPU計(jì)算后返回結(jié)果，再把結(jié)果送給GPU再次計(jì)算，可以看到，大量的時(shí)間都耗費(fèi)在數(shù)據(jù)傳輸上了。

稍作調(diào)整之后，當(dāng)預(yù)先把指令都放入計(jì)算時(shí)，內(nèi)存的傳輸降為一次即可完成相同的任務(wù)。

如果換成pyTorch的代碼就是把兩行代碼轉(zhuǎn)為一行x.cos().cos()，效率能夠提升兩倍。

不過這種優(yōu)化措施并不是在所有場景下都適用。因?yàn)镚PU預(yù)先需要知道所有執(zhí)行的指令，并生成CUDA代碼，所以無法在eager-mode下使用。而且并非所有的運(yùn)算符融合都像pointwise操作符這么簡單。

如果你曾經(jīng)寫過CUDA內(nèi)核代碼的話，就可以知道任何兩個PyTorch都有機(jī)會進(jìn)行融合來節(jié)省全局內(nèi)存的讀寫成本。現(xiàn)有的編譯器如NVFuser和XLA通常只能進(jìn)行一些簡單的融合，肯定比不上AI工程師的設(shè)計(jì)。如果你想嘗試自己編寫一些定制的CUDA內(nèi)核，Triton就比較適合新手入門。

運(yùn)算符融合的效果就是更多的操作，時(shí)間成本相同，這也是為什么激活函數(shù)的計(jì)算成本幾乎都是一樣的，盡管gelu顯然比relu多了很多操作。

當(dāng)需要推理你的操作是否有內(nèi)存帶寬限制時(shí)，calculator可以發(fā)揮很大的作用。

對于簡單的算子來說，可以直接推理內(nèi)存帶寬。例如，A100有1.5T字節(jié)/秒的全局內(nèi)存帶寬，可以進(jìn)行19.5T FLOPS的計(jì)算。因此，如果你使用32位浮點(diǎn)（即4個字節(jié)），GPU可以執(zhí)行20萬億次操作的相同時(shí)間內(nèi)加載4000億個數(shù)字。此外，為了執(zhí)行一個簡單的單項(xiàng)運(yùn)算（如把一個tensor乘2），實(shí)際上需要將tensor寫回全局內(nèi)存。所以將單項(xiàng)運(yùn)算做了大約一百次以后，才能夠等到內(nèi)存數(shù)據(jù)送進(jìn)來。

在像NVFuser這樣的融合編譯器的幫助下，實(shí)際上可以很容易地測量成本。

以一個PyTorch函數(shù)為例，并用融合編譯器對其進(jìn)行基準(zhǔn)測試，然后就可以計(jì)算出不同的重復(fù)值所達(dá)到的FLOPS和內(nèi)存帶寬。

增加重復(fù)次數(shù)是在不增加內(nèi)存訪問的情況下增加計(jì)算量的一個簡單方法，這也被稱為增加計(jì)算強(qiáng)度。

因?yàn)閠ensor的大小為N，需要將執(zhí)行2*N次內(nèi)存訪問，以及N*repeat FLOP。因此，實(shí)現(xiàn)的內(nèi)存帶寬將是byte_per_elem * 2 * N / itrs_per_second，而實(shí)現(xiàn)的FLOPS將是N * repeat / itrs_per_second。

把運(yùn)行時(shí)間、flops和實(shí)現(xiàn)的內(nèi)存帶寬取對數(shù)后繪制出來的結(jié)果可以看到，執(zhí)行64次乘法之前，運(yùn)行時(shí)間并沒有明顯的增加。這也意味著，在這之前，內(nèi)存帶寬是有限的，計(jì)算大部分是閑置的。

因此，一開始只實(shí)現(xiàn)了0.2 teraflops。當(dāng)我們把計(jì)算強(qiáng)度提高一倍時(shí)，這個數(shù)字就會線性增長，直到我們接近9.75 teraflops的峰值，也就是「計(jì)算極限」。

內(nèi)存帶寬開始時(shí)接近峰值，隨著計(jì)算強(qiáng)度的增加，開始下降。這也符合預(yù)期，因?yàn)閷?shí)際上更多的時(shí)間花在了實(shí)際的計(jì)算上，而非訪問內(nèi)存。

在這種情況下可以很容易看到什么時(shí)候是計(jì)算約束，什么時(shí)候是內(nèi)存約束。

對于重復(fù)次數(shù)小于32次時(shí)，內(nèi)存帶寬已經(jīng)飽和，而計(jì)算能力卻沒有得到充分利用。相反，一旦重復(fù)大于64次，會發(fā)現(xiàn)計(jì)算量已經(jīng)飽和（即達(dá)到接近峰值FLOPS），而內(nèi)存帶寬利用率開始下降。

對于更大的系統(tǒng)，通常很難說是計(jì)算約束還是內(nèi)存帶寬約束，因?yàn)榭赡芡瑫r(shí)包含了計(jì)算約束和內(nèi)存約束。

衡量計(jì)算約束程度的一個常見方法是，將你實(shí)現(xiàn)的FLOPS作為峰值FLOPS的一個百分比作為指標(biāo)。如果實(shí)現(xiàn)了峰值FLOPS的80%，那就說明計(jì)算資源利用的比較充分，其余的時(shí)間可能是花在內(nèi)存帶寬上了。

其他開銷

代碼中沒有花在傳輸或計(jì)算tensor的時(shí)間都稱為開銷（overhead），比如花在Python解釋器上的時(shí)間，花在PyTorch框架上的時(shí)間，花在啟動CUDA內(nèi)核（但不執(zhí)行）的時(shí)間都是開銷。

開銷之所以成為一個問題，主要原因是現(xiàn)代GPU的速度非?？臁Ｒ粋€A100可以每秒進(jìn)行312萬億次的浮點(diǎn)運(yùn)算（312 TeraFLOPS）。相比之下，Python的運(yùn)行速度就相當(dāng)慢了，一秒鐘內(nèi)只能進(jìn)行3200萬次加法運(yùn)算。

這也意味著，在Python可以執(zhí)行一個FLOP的時(shí)間里，A100可以運(yùn)行975萬FLOPS。

像PyTorch這樣的框架在進(jìn)入實(shí)際內(nèi)核之前也有很多層調(diào)度。如果你用PyTorch做同樣的實(shí)驗(yàn)，每秒只能得到28萬次操作。當(dāng)然，執(zhí)行小tensor并不是建立PyTorch的目的，但是如果確實(shí)在科學(xué)計(jì)算中使用小tensor，你就會發(fā)現(xiàn)PyTorch與C++相比慢得驚人。

一個更直觀的圖可以看到，PyTorch執(zhí)行一個加法時(shí)產(chǎn)生的配置文件，除了一個小方塊外，其余所有的都是純開銷。

現(xiàn)代深度學(xué)習(xí)模型通常都在進(jìn)行大規(guī)模的計(jì)算操作，并且像PyTorch這樣的框架是異步執(zhí)行的。也就是說，當(dāng)PyTorch正在運(yùn)行一個CUDA內(nèi)核時(shí)，它可以繼續(xù)運(yùn)行并在后面排起更多的CUDA內(nèi)核。因此，只要PyTorch能夠「提前」運(yùn)行CUDA內(nèi)核，大部分的框架開銷就會被完全隱藏起來