科技獵機(jī)器學(xué)習(xí)陷入困境！谷歌大腦專家發(fā)文吐槽AI工程現(xiàn)狀
論文作者 | Paul Barham，Michael Isard
編譯 | Maglish
編輯 | 陳思
AI 前線導(dǎo)讀：數(shù)值計(jì)算系統(tǒng)的性能和可編程性正陷入低谷。系統(tǒng)研究人員出色的工作使機(jī)器學(xué)習(xí)的基準(zhǔn)在過去 5 年不斷提升，但是探索創(chuàng)新的機(jī)器學(xué)習(xí)研究想法變得越來越難。在本文中，Google Brain 的研究人員解釋了硬件加速器的進(jìn)化如何有利于編譯器后端優(yōu)化巨大的單內(nèi)核，介紹了這種對高性能但不靈活的內(nèi)核的依賴如何增強(qiáng)了編程模型的主導(dǎo)風(fēng)格。作者認(rèn)為這些抽象編程缺乏表現(xiàn)性、可維護(hù)性和模塊性，而這些都阻礙了研究的進(jìn)展。最后，作者指出了該領(lǐng)域的發(fā)展方向，并提倡在現(xiàn)代加速器上推進(jìn)高性能通用數(shù)值計(jì)算系統(tǒng)的發(fā)展。
更多優(yōu)質(zhì)內(nèi)容請關(guān)注微信公眾號“AI 前線”（ID：ai-front）
在試圖改進(jìn)膠囊網(wǎng)絡(luò)的實(shí)現(xiàn)，以擴(kuò)大到更大的數(shù)據(jù)集時(shí)，研究團(tuán)隊(duì)有了這篇論文的初步想法。膠囊網(wǎng)絡(luò)是一個(gè)令人興奮的機(jī)器學(xué)習(xí)研究思想，其中標(biāo)量值的“神經(jīng)元”被小矩陣取代，使它們能夠捕捉更復(fù)雜的關(guān)系。膠囊或許不是機(jī)器學(xué)習(xí)中的“下一個(gè)大事件”，但它算是創(chuàng)新性機(jī)器學(xué)習(xí)研究理念的一個(gè)代表性例子。
雖然卷積膠囊模型需要的浮點(diǎn)運(yùn)算（FLOPs）比卷積神經(jīng)網(wǎng)絡(luò)小 4 倍，并且訓(xùn)練參數(shù)少 16 倍，但是用 TensorFlow 和 PyTorch 實(shí)現(xiàn)的膠囊模型都比卷積神經(jīng)網(wǎng)絡(luò)模型慢得多，并且很小的模型就已耗盡了內(nèi)存。到底是為什么呢？
膠囊網(wǎng)絡(luò)內(nèi)部循環(huán)的簡化形式類似于傳統(tǒng)的 CNN 層中的計(jì)算，但是是對 4×4 的矩陣進(jìn)行操作，而不是標(biāo)量。
當(dāng)前機(jī)器學(xué)習(xí)框架的一個(gè)基本組成部分是 2D 卷積。大多數(shù)框架提供一個(gè)原語操作，輸入大小為 H×W 的 N 個(gè)圖像，其中每個(gè)像素的“深度”為 Ci 個(gè)通道。對于“核大小”k＝3 和“步長”s＝2，conv2D 計(jì)算以每個(gè)（x，y）坐標(biāo)為中心的重疊的 3×3 像素塊的加權(quán)和，然后產(chǎn)生像素深度為 Co 的 N 個(gè)較小的圖像（如圖 1 所示）。數(shù)學(xué)上可以表達(dá)如下：
其中的點(diǎn)乘表示標(biāo)量乘法，O,I 和 K 都是 4 維的標(biāo)量矩陣。最后的代碼只是圍繞一個(gè)乘法累加操作的 7 個(gè)嵌套循環(huán)，但矩陣布局、矢量化、并行和緩存對最后的表現(xiàn)都很重要。
卷積膠囊計(jì)算 3×3 卷積塊中“姿態(tài)”矩陣的加權(quán)和，形成“投票”：
其中的點(diǎn)乘表示矩陣乘法，V,P 和 W 是 4×4 矩陣的 4 維數(shù)組，或等價(jià)于 6 維標(biāo)量矩陣。
下面將解釋為什么機(jī)器學(xué)習(xí)框架難以有效地運(yùn)行膠囊計(jì)算。
卷積膠囊的原語可以在 CPU 上合理地實(shí)現(xiàn)（表 1），但是在加速器（例如 GPU 和 TPU）上就會出現(xiàn)問題。在加速器上的表現(xiàn)至關(guān)重要，因?yàn)閹缀醍?dāng)前所有的機(jī)器學(xué)習(xí)研究，以及大多數(shù)生產(chǎn)模型的訓(xùn)練，都需要使用加速器。使用加速器執(zhí)行特定的 ML 訓(xùn)練或大量推理工作的邊際成本比使用 CPU 要低得多。
加速器非常適用于機(jī)器學(xué)習(xí)，因?yàn)橛?xùn)練任務(wù)中計(jì)算量較大的部分是密集的多維數(shù)組的線性代數(shù)運(yùn)算。密集的線性代數(shù)與 CPU 所承擔(dān)的工作負(fù)載相比是比較規(guī)則的，并且相對容易并行化。因此，人們?yōu)槌Ｒ?guī)并行計(jì)算設(shè)計(jì)了越來越復(fù)雜的加速器。常見的加速器特性包括“warp”、塊和線程網(wǎng)格、非常寬的矢量計(jì)算單元（ALU）和脈動陣列乘法器（MXUs）。但即使這些常規(guī)計(jì)算也很難獲得良好的加速器性能。雖然常見的計(jì)算方式受到關(guān)注，并得到了很好的優(yōu)化，但是非標(biāo)準(zhǔn)計(jì)算的性能會受影響，如卷積膠囊。
很難從常規(guī)計(jì)算中獲得良好性能的一個(gè)主要原因是編譯器必須考慮加速器的內(nèi)存系統(tǒng)以及 ALU。為了防止數(shù)據(jù)瓶頸，加速器的并行能力已經(jīng)與內(nèi)存系統(tǒng)緊密耦合。GPU 上的 ALU 性能峰值需要“聯(lián)合負(fù)載”，一個(gè) warp 的所有 32 個(gè)線程同時(shí)訪問同一緩存線中的不同值；程序?qū)崿F(xiàn)必須適應(yīng)內(nèi)存的大小和步長；并且高效的程序必須利用單次內(nèi)存訪問中加載的所有值。
通常，加速器代碼必須通過內(nèi)存結(jié)構(gòu)執(zhí)行顯式調(diào)度，而不是依賴透明的多級緩存。內(nèi)存訪問粒度一般要求線程協(xié)同加載彼此的值，然后交換它們；這樣，代碼還包含跨循環(huán)迭代的復(fù)雜指令調(diào)度。雖然匹配并行 ALU 的內(nèi)存訪問能帶來良好的硬件利用率，但任何不匹配都會導(dǎo)致性能下降一個(gè)數(shù)量級。為了避免這種情況，需要為每一代加速器調(diào)整內(nèi)核參數(shù)，例如 padding、步長和維度布局。
對于像卷積這樣的“模版計(jì)算”，輸入值需要被重疊的計(jì)算窗口重復(fù)使用，調(diào)度負(fù)載和存儲以優(yōu)化存儲器帶寬是非常具有挑戰(zhàn)性的任務(wù)。而卷積膠囊中的數(shù)據(jù)重用模式具有幾個(gè)額外的復(fù)雜性維度。
由于對參數(shù)調(diào)整分析很困難，加速器的高性能后端在一組計(jì)算“內(nèi)核”（通常是孤立的循環(huán)嵌套）上花費(fèi)了大量的開發(fā)工作，例如 2D 卷積和批矩陣乘法，它們決定了基準(zhǔn)的表現(xiàn)性能。對于這些內(nèi)核中的每一個(gè)，后端維護(hù)人員需要花費(fèi)數(shù)小時(shí)或數(shù)天為一小組具有代表性的操作數(shù)形狀（矩陣維度的基數(shù)）尋找最佳算法和參數(shù)設(shè)置，然后在運(yùn)行時(shí)使用啟發(fā)式或自動調(diào)整選擇這些預(yù)調(diào)整的實(shí)現(xiàn)。
對于機(jī)器學(xué)習(xí)論文來說，提出新原語，卻無法用現(xiàn)有內(nèi)核有效計(jì)算的情況是非常常見的。研究人員開發(fā)了像張量理解（Tensor Comprehensions，TC）和 PlaidML 這樣的編譯器，允許終端用戶編寫自定義內(nèi)核，它們都為領(lǐng)域特定語言提供了與數(shù)學(xué)公式相似的簡潔語法，例如圖 2 中膠囊原語的 TC 實(shí)現(xiàn)與公式 2 的對比。
但現(xiàn)實(shí)情況是，這些工具只適用于編譯小代碼片段。編譯時(shí)間通常很長，并且代碼質(zhì)量往往不能達(dá)到峰值性能。
圖 3a 說明了當(dāng)前編譯器框架為傳統(tǒng) 2D 卷積生成 GPU 代碼的困難，其性能需要與 cuDNN 中經(jīng)過精心調(diào)整的庫競爭。在每一種情況下，作者使用框架中可用的最低級別原語實(shí)現(xiàn)了 conv2d（公式 1）。TC 使用遺傳搜索算法來優(yōu)化參數(shù)，性能僅為 cuDNN 的八分之一，但是僅用了一個(gè)小時(shí)的搜索時(shí)間。最終性能也非常依賴于輸入操作數(shù)（NCHW vs NHWC）的內(nèi)存布局。TVM 具有類似的自動調(diào)整卷積模板，在 30 分鐘搜索之后亦無法與 cuDNN 的性能相抗衡（圖 3b）。PlaidML 獲得最佳性能，比 cuDNN 慢 4 倍，編譯時(shí)間快。TVM 也為這些操作數(shù)形狀提供了手工調(diào)度的 conv2d 內(nèi)核，但是它比 cuDNN 慢 19 倍。
回到膠囊示例中，作者接下來嘗試為核心的膠囊原語編寫自定義內(nèi)核（公式 2）。用 gcc 圍繞 4×4 矩陣乘法（matmul）函數(shù)編譯明顯的 C++ 循環(huán)嵌套，可以得到了高質(zhì)量的矢量化代碼，將其作為基線。它在一個(gè) x86 內(nèi)核上運(yùn)行約 60ms，用 OpenMP 在 6 個(gè)內(nèi)核并行化時(shí)達(dá)到 11.7ms。自己編寫的 CUDA 實(shí)現(xiàn)運(yùn)行了 1.9ms，但花費(fèi)了兩天時(shí)間進(jìn)行手動調(diào)整。
雖然 PlaidML 在 gcc 上編譯得很快，但內(nèi)核執(zhí)行要慢得多。TC 需要近 3 分鐘來找到一個(gè)優(yōu)于 CPU 的內(nèi)核，但最終發(fā)現(xiàn)了運(yùn)行時(shí)間少于 1.8ms 的調(diào)度（見表 1 和圖 3C）。
作者對這些結(jié)果給出如下解釋：當(dāng)前框架在工作負(fù)載上較為先進(jìn)，手動調(diào)整由特定模型或模型簇使用的小計(jì)算集是合理的。不幸的是，該框架并不適合研究，因?yàn)閷Ψ浅Ｒ?guī)計(jì)算做實(shí)驗(yàn)時(shí)存在性能懸崖。雖然在生產(chǎn)部署前花費(fèi)幾個(gè)小時(shí)的搜索時(shí)間是可以接受的，但期望研究人員忍受這樣的編譯時(shí)間是不現(xiàn)實(shí)的（要記住這只是一個(gè)內(nèi)核，而需要編譯的可能是一個(gè)巨大的整體計(jì)算）；即使優(yōu)化的內(nèi)核定期緩存在本地，它也是傳播研究的主要障礙。任何下載模型源代碼的人都必須花費(fèi)數(shù)小時(shí)或幾天的時(shí)間在硬件上進(jìn)行編譯，然后才能進(jìn)行實(shí)驗(yàn)。
在像 TensorFlow 和 PyTorch 這樣的 ML 框架中使用自定義計(jì)算并不直接。因此，在 TensorFlow 和 PyTorch 中實(shí)現(xiàn)卷積膠囊的最簡單和最有效的方法是使用已經(jīng)由這些框架支持的高級操作。
不幸的是，兩個(gè)框架都不能將最終的縮減融合到批處理矩陣乘法（預(yù)優(yōu)化內(nèi)核）中，這大大增加了所需的存儲器帶寬以及中間存儲要求。為了計(jì)算兩個(gè)相對較小的量，API 迫使程序?qū)⒈日嬲钄?shù)據(jù)高出兩個(gè)數(shù)量級的數(shù)據(jù)復(fù)制、重新排列和實(shí)現(xiàn)到內(nèi)存。正是上述問題讓研究人員無法找到任何高性能的卷積膠囊的實(shí)現(xiàn)。
即使對于像 ResNet 那樣簡單的 ML 模型，操作數(shù)也有不同的形狀，因此不同的卷積調(diào)用可能具有不同的最優(yōu)參數(shù)。如果中間值的生成操作和消費(fèi)操作之間的布局不同，那么該值必須被“轉(zhuǎn)置”（轉(zhuǎn)換為不同的布局）。
ML 計(jì)算與許多計(jì)算不同，因?yàn)樗鼈兺ǔＩ婕敖频母↑c(diǎn)數(shù)或定點(diǎn)數(shù)。通過改變精度或非均勻量化的選擇，可以或多或少地獲得目標(biāo)度量，例如測試集預(yù)測精度。精度也會影響計(jì)算 / 通信瓶頸。
常見子表達(dá)式消除（CSE）： CSE 對于 ML 框架非常重要，因?yàn)樗捎昧朔聪騻鞑ビ?jì)算結(jié)構(gòu)。標(biāo)準(zhǔn) CSE 傾向于將激活值物化到內(nèi)存，但是加速器內(nèi)存的數(shù)量有限以及 ALU 的緩慢，意味著人們更傾向于重新計(jì)算激活。因此，CSE 為框架引入了另一個(gè)組合搜索問題：選擇哪些值應(yīng)該被物化。
分布式執(zhí)行： 由于數(shù)據(jù)并行結(jié)構(gòu)，機(jī)器學(xué)習(xí)計(jì)算通常可以有效地分布在多個(gè)加速器上。在實(shí)踐中，編譯器和機(jī)器學(xué)習(xí)框架通常利用與設(shè)備中使用的機(jī)制不相交的方法來實(shí)現(xiàn)分布式并行性。據(jù)我們所知，沒有框架試圖共同優(yōu)化選擇計(jì)算片段應(yīng)該在哪個(gè)設(shè)備上運(yùn)行，以及如何選擇子計(jì)算在設(shè)備內(nèi)的結(jié)構(gòu)。
機(jī)器學(xué)習(xí)訓(xùn)練算法特別依賴于自動優(yōu)化策略，因?yàn)榇蠖鄶?shù)計(jì)算通常在由源代碼中出現(xiàn)的“前向傳遞”中合成的梯度算子中執(zhí)行。由于計(jì)算梯度的代碼不是程序員編寫的，因此程序員來引導(dǎo)優(yōu)化的機(jī)會有限。即使在沒有自動生成梯度的情況下，也很難寫出有手動優(yōu)化標(biāo)注的模塊化代碼。
最近的研究對自動全程序優(yōu)化技術(shù)的興趣越來越大，但是方法還在起步階段，并且通常集中每次只優(yōu)化程序的一個(gè)方面。毫無疑問，多維整體程序優(yōu)化是一項(xiàng)艱巨的任務(wù)，但或許可以從最近混合搜索、學(xué)習(xí)方法的成功中獲得一些希望，如 AlphaGo，它表明在巨大的組合搜索空間中有希望找到好的解決方案。似乎有必要在搭建機(jī)器學(xué)習(xí)框架時(shí)考慮自動優(yōu)化器，才可能充分利用全程序優(yōu)化。
作者首先觀察到數(shù)值計(jì)算得益于傳統(tǒng)編程語言中不存在的特性。自動微分就是一個(gè)這樣的特性，并且數(shù)值程序也是不尋常的，因?yàn)樗鼈兪褂迷谄鋮?shù)秩（維數(shù)）上具有多態(tài)性的函數(shù)進(jìn)行編寫。再考慮標(biāo)準(zhǔn)卷積表達(dá)式（公式 1）。對批（n）和輸出通道（Co）維度的每個(gè)元素的計(jì)算是獨(dú)立的，并且表達(dá)卷積的自然方式是以三維輸入 I 和 K 所寫的子計(jì)算。
當(dāng)輸出具有更多維度時(shí)，語言可以自動地“提升”函數(shù)，跨越批和輸出通道維度。早在 APL 語言中，數(shù)字語言就已經(jīng)包含了提升秩多態(tài)性，但是關(guān)于如何將這種多態(tài)性與現(xiàn)代模塊化類型和語言結(jié)合起來，還存在大量的開放性研究問題。
回想一下，后端是圍繞對大型單內(nèi)核的調(diào)用而構(gòu)建的。作者認(rèn)為這種后端設(shè)計(jì)方法正在阻礙編程模型的可維護(hù)性、可調(diào)試性和表達(dá)性方面的進(jìn)展。更糟糕的是，由此帶來的對語言創(chuàng)新的阻礙本身就減少了后端開發(fā)人員對當(dāng)前情況進(jìn)行改進(jìn)的動力。
單后端內(nèi)核的一個(gè)結(jié)果是前端選擇內(nèi)核或“操作符”作為抽象點(diǎn)。在流行的框架中，如 TensorFlow 和 PyTorch，用戶程序是在 Python 中編寫的，并調(diào)用以后端特定的語言和庫（如 C++、MKL、CUDA 和 cuDNN）編寫的運(yùn)算符，或者有時(shí)是較低級別但可移植的特定領(lǐng)域語言，例如 Tile 和張量理解。當(dāng)現(xiàn)有的運(yùn)算符不足以執(zhí)行任務(wù)時(shí)，用戶必須使用較低級別的語言來編寫新的運(yùn)算符，通常還必須手動編寫其梯度，這是一個(gè)困難且容易出錯(cuò)的過程。
有的框架，例如 Julia，它名義上使用相同的語言來表示操作符圖及其實(shí)現(xiàn)，但是后端設(shè)計(jì)抵消了前端的有效性。
正如先前所建議的，語言層面的另一種選擇是提供低維 conv2d 函數(shù)，并在必要時(shí)將其提升到更高的維度。作者認(rèn)為這將大大提高可讀性和可維護(hù)性。單后端內(nèi)核是這種方法的一個(gè)障礙，因?yàn)檎Z言必須自動發(fā)現(xiàn)在某些維度排序下等價(jià)于單內(nèi)核的模式，然后在調(diào)用內(nèi)核之前自動重新排序維度。這種模式匹配如果出現(xiàn)任何失敗，都會使性能嚴(yán)重下降。但是，如果有個(gè)編譯器可以為通用提升函數(shù)編譯出高效代碼，那么進(jìn)程會快得多。
人們經(jīng)常提倡“命名維度”，即讓矩陣的維數(shù)與文本名稱相關(guān)聯(lián)，或者與其數(shù)字位置相關(guān)聯(lián)。命名維度能夠提高可讀性，但作者認(rèn)為，在提升代碼模塊性方面它們的影響可能會更大，因?yàn)槊S度可以使語言遠(yuǎn)離固定維度順序，從而讓函數(shù)提升更方便。例如在 APL 中，維度有嚴(yán)格的順序，在調(diào)用函數(shù)前參數(shù)必須經(jīng)過轉(zhuǎn)置才能確保正確的順序。
為了有效地實(shí)現(xiàn)具有命名或無序維度的程序，有必要重新考慮后端的設(shè)計(jì)。有幾個(gè)項(xiàng)目已經(jīng)在這個(gè)方向上采取了初步嘗試。他們在一定程度上解相關(guān)源代碼的維數(shù)順序與低級代碼的維度順序，但是仍然要求前端指定每個(gè)陣列的有序維度。
目前針對上述的這些問題，還沒有很好的解決方案。一方面是工程實(shí)現(xiàn)上存在難度，一方面是沒有足夠的動機(jī)去研究這些問題的解決方案。面向機(jī)器學(xué)習(xí)的端到端工具鏈的研究需要許多不同領(lǐng)域的專業(yè)知識。很遺憾的是，盡管研究人員在一些子問題上已經(jīng)做出了效果卓越的工作或者一些具有前瞻性的努力，但是仍無法提出一個(gè)能基本解決端到端問題的綜合性思路。本著逐個(gè)擊破的想法，作者提出了一些探索性的思路：
- 語言設(shè)計(jì)，包括自動化微分以及使用語言的語法表達(dá)純命名維度和內(nèi)核。
- 定義布局不可知的通用循環(huán)嵌套圖的后端 IR(接口需求，Interface Requirements)。
- 轉(zhuǎn)化上述 IR，將其解構(gòu)為具體的 CSE 策略。
- 在上面提到的降低 IR 的情況下生成加速器代碼的編譯過程，快速產(chǎn)生足夠的代碼并接近最好的表現(xiàn)。
對于現(xiàn)有的機(jī)器學(xué)習(xí)工具鏈，這篇論文并不是輕視他們的思想和工程實(shí)現(xiàn)，并且顯而易見的是它們對于大部分領(lǐng)域相關(guān)人員來說是非常具有價(jià)值的。作者的考慮主要來自于語言和后端的不靈活性會成為阻礙創(chuàng)新型研究的真正障礙，這會使得這個(gè)非常活躍的領(lǐng)域進(jìn)展緩慢。因此在設(shè)計(jì)加速器、工具鏈時(shí)，作者希望大家能夠更多的考慮這一點(diǎn)，尤其是對于基準(zhǔn)測試的研究。
查看論文原文：
https://dl.acm.org/citation.cfm?id=3321441
