在圖像、語音識別、自然語言處理、強(qiáng)化學(xué)習(xí)等許多技術(shù)領(lǐng)域中，深度學(xué)習(xí)已經(jīng)被證明是非常有效的，并且在某些問題上已經(jīng)達(dá)到甚至超越了人類的水平。然而，深度學(xué)習(xí)對于計(jì)算能力有著很大的依賴，除了改變模型和算法，是否可以從系統(tǒng)的層面來優(yōu)化深度學(xué)習(xí)計(jì)算，進(jìn)而改善計(jì)算資源的使用效率?本文中，來自微軟亞洲研究院異構(gòu)計(jì)算組資深研究員伍鳴與大家分享他對深度學(xué)習(xí)計(jì)算優(yōu)化的一些看法。
深度學(xué)習(xí)在近幾年里取得了巨大的進(jìn)步，它已經(jīng)或者是有望成功地被應(yīng)用在我們許多生活場景中，比如自動(dòng)駕駛、安防、翻譯、醫(yī)療等等�？梢哉f，計(jì)算機(jī)的計(jì)算和通信能力的大幅提升是促使深度學(xué)習(xí)成功的重要因素。
深度學(xué)習(xí)為什么依賴于超大的計(jì)算能力?
首先，深度學(xué)習(xí)本質(zhì)上是基于統(tǒng)計(jì)的科學(xué)，所以大規(guī)模的樣本數(shù)據(jù)對于深度學(xué)習(xí)的效果是至關(guān)重要的。其次，更大規(guī)模和更復(fù)雜的神經(jīng)網(wǎng)絡(luò)模型已經(jīng)被證明非常有效，并在產(chǎn)品中有廣泛的使用，這同時(shí)也產(chǎn)生了對計(jì)算能力的更大要求和消耗。舉個(gè)例子，具有8層神經(jīng)元的AlexNet網(wǎng)絡(luò)2012年在ImageNet數(shù)據(jù)集上取得16%的錯(cuò)誤率，該網(wǎng)絡(luò)的一次迭代運(yùn)行大約需要1.4 GFLOP的計(jì)算量。而微軟提出的使用152層神經(jīng)元的殘差網(wǎng)絡(luò)(ResNet)于2015年在該數(shù)據(jù)集上取得3.5%的錯(cuò)誤率，其一次迭代的計(jì)算量大約是22.6GFLOP，是AlexNet的16倍。在當(dāng)今的生產(chǎn)環(huán)境中，圖像、語音以及自然語言處理相關(guān)的模型，例如人臉識別、語音轉(zhuǎn)文字、機(jī)器翻譯等，即使給予相當(dāng)多的計(jì)算資源，很多仍需要幾周的時(shí)間才能完成訓(xùn)練。
再次，深度學(xué)習(xí)模型是迅速迭代的。在AI領(lǐng)域，每年學(xué)術(shù)界和工業(yè)界都會(huì)提出大量的新模型。對每一個(gè)實(shí)際的問題，開發(fā)者需要不斷嘗試不同的模型和算法，甚至對于同一種模型算法，也需要去反復(fù)調(diào)試超參數(shù)以獲得最好的預(yù)測效果�？上攵�，如果模型的每次訓(xùn)練都要幾周的時(shí)間，那么尋找最優(yōu)模型的過程會(huì)非常漫長和痛苦。
另外，模型的線上推理具有更加極致的性能要求。線上的服務(wù)具有硬性的服務(wù)等級協(xié)議(SLA)，所以在實(shí)際部署大型模型時(shí)，需要手工重新優(yōu)化在深度學(xué)習(xí)框架(如TensorFlow)上已經(jīng)訓(xùn)練好的模型，導(dǎo)致大量額外工程開銷的產(chǎn)生。
由此可見，進(jìn)一步優(yōu)化深度學(xué)習(xí)計(jì)算對于深度學(xué)習(xí)的快速發(fā)展和成功應(yīng)用起著至關(guān)重要的作用。
深度學(xué)習(xí)計(jì)算優(yōu)化的挑戰(zhàn)和機(jī)會(huì)
目前，優(yōu)化深度學(xué)習(xí)的計(jì)算存在以下幾個(gè)主要的挑戰(zhàn)：
1)單機(jī)單計(jì)算單元(如GPU)的資源限制往往不能滿足對大規(guī)模數(shù)據(jù)和模型的處理要求，那么就需要使用多機(jī)多計(jì)算單元來橫向擴(kuò)展計(jì)算的規(guī)模。如何才能最大限度地減少通信的開銷從而最大化多機(jī)的并行度?
2)如何優(yōu)化神經(jīng)網(wǎng)絡(luò)的計(jì)算使得它能夠把單個(gè)硬件計(jì)算單元的效率發(fā)揮到極致?
3)雖然許多硬件計(jì)算單元(GPU、FPGA等)的計(jì)算能力很強(qiáng)大，但是它們的內(nèi)存資源(即設(shè)備內(nèi)存)非常稀缺。當(dāng)它們不能提供模型運(yùn)行所需要的內(nèi)存資源時(shí)，要么運(yùn)算不能夠進(jìn)行下去，要么就需要將計(jì)算所需的數(shù)據(jù)在主存和設(shè)備內(nèi)存之間倒來倒去，帶來很大的運(yùn)行開銷。如何才能更好地利用有限的設(shè)備內(nèi)存資源從而不給計(jì)算效率帶來負(fù)面的影響?
4)深度學(xué)習(xí)開發(fā)者和研究人員通常只想關(guān)注神經(jīng)網(wǎng)絡(luò)模型和算法本身，并不想被復(fù)雜的優(yōu)化問題分散精力。這意味著深度學(xué)習(xí)框架這樣的系統(tǒng)軟件最好能夠?qū)崿F(xiàn)自動(dòng)優(yōu)化，而對模型開發(fā)者透明。那么，如何對特定的優(yōu)化做合理的抽象使其更加靈活通用、更加容易地集成在系統(tǒng)框架中便是需要認(rèn)真考慮的問題。
事實(shí)上，任何方面的優(yōu)化問題都可以從模型算法和系統(tǒng)兩個(gè)角度來看待。一方面，我們可以通過改變模型和算法來優(yōu)化其對計(jì)算資源的使用效率從而改進(jìn)其運(yùn)行速度。這樣的優(yōu)化對特定的算法往往非常有效，但卻不容易擴(kuò)展應(yīng)用到其它算法中。而另一方面，也就是微軟亞洲研究院異構(gòu)計(jì)算組正在進(jìn)行的研究，則是在系統(tǒng)中實(shí)施模型算法無關(guān)的優(yōu)化，這樣的優(yōu)化，通�？梢詾楦嗟膽�(yīng)用帶來性能的好處，同時(shí)也符合我們在前文提到的透明性的要求。
以系統(tǒng)優(yōu)化助力深度學(xué)習(xí)計(jì)算
為了能夠更好地理解系統(tǒng)這一層面的優(yōu)化，我們先來簡單介紹一下深度學(xué)習(xí)框架系統(tǒng)的背景知識。當(dāng)今工業(yè)界流行的深度學(xué)習(xí)系統(tǒng)(包括TensorFlow、PyTorch、CNTK、MxNet、Caffe等)大都采用分層的體系結(jié)構(gòu)設(shè)計(jì)。在前端提供高級語言(例如Python)的接口抽象，允許用戶方便地描述神經(jīng)網(wǎng)絡(luò)結(jié)構(gòu)，也就是深度學(xué)習(xí)的模型。描述好的模型在被系統(tǒng)運(yùn)行前，首先會(huì)被轉(zhuǎn)換成數(shù)據(jù)流圖(Data-flow Graph)。在這個(gè)數(shù)據(jù)流圖中，節(jié)點(diǎn)是特定的矩陣操作(也就是Operator，如Sigmoid、Matrix Multiplication等)，而連接不同節(jié)點(diǎn)的邊則是操作節(jié)點(diǎn)的輸入和輸出矩陣。這個(gè)數(shù)據(jù)流圖也可以被看成是深度學(xué)習(xí)計(jì)算的中間表達(dá)。然后，深度學(xué)習(xí)系統(tǒng)的后端將這個(gè)數(shù)據(jù)流圖映射到實(shí)際硬件上進(jìn)行高效地執(zhí)行，而大部分系統(tǒng)層面的優(yōu)化就是在這個(gè)階段完成的。
加速分布式深度學(xué)習(xí)訓(xùn)練
分布式訓(xùn)練的主要瓶頸在于多機(jī)之間的通信開銷。如今計(jì)算機(jī)網(wǎng)絡(luò)的硬件技術(shù)已經(jīng)有了很大的發(fā)展，InfiniBand的RDMA網(wǎng)卡(Remote Direct Memory Access，這是一種硬件的網(wǎng)絡(luò)技術(shù)，它使得計(jì)算機(jī)訪問遠(yuǎn)程的內(nèi)存時(shí)無需遠(yuǎn)程機(jī)器上CPU的干預(yù))已經(jīng)可以提供50~100Gbps的網(wǎng)絡(luò)帶寬和微秒級的傳輸延遲。目前許多以深度學(xué)習(xí)為目標(biāo)應(yīng)用的GPU機(jī)群都部署了這樣的網(wǎng)絡(luò)。然而深度學(xué)習(xí)系統(tǒng)如何才能充分利用好硬件提供的通信能力使分布式的訓(xùn)練獲得更大的性能提升呢?另外，使用RDMA的軟件接口進(jìn)行通信能夠繞過TCP/IP協(xié)議棧，減少了操作系統(tǒng)內(nèi)核態(tài)的運(yùn)行開銷。在這樣的網(wǎng)絡(luò)通信技術(shù)的支持下，任何與通信相關(guān)的計(jì)算處理的開銷都會(huì)變得非常顯著，而這正是許多原先基于TCP/IP而設(shè)計(jì)的網(wǎng)絡(luò)通信機(jī)制中所存在的問題。
RPC(Remote Procedure Call，遠(yuǎn)程過程調(diào)用)是一個(gè)被廣泛使用的多機(jī)之間的通信抽象原語，它的主要設(shè)計(jì)目標(biāo)是通用性。在沒有考慮RDMA的情況下，很多深度學(xué)習(xí)框架都會(huì)采用RPC的機(jī)制(例如gRPC)來實(shí)現(xiàn)多機(jī)之間的通信。然而，RPC需要維護(hù)一個(gè)內(nèi)部的私有緩存，從而不得不引入用戶數(shù)據(jù)的存儲空間和內(nèi)部緩存之間的數(shù)據(jù)拷貝。這種內(nèi)存拷貝的開銷在使用RDMA網(wǎng)絡(luò)的情況下會(huì)變得非常明顯。我們通過micro-benchmark觀察到，跟使用基于TCP/IP的gRPC相比，直接通過RDMA的接口傳輸消息(對不同的消息大小)可以有2到10倍的性能提升。
那么針對深度學(xué)習(xí)的應(yīng)用負(fù)載，如何才能更好地利用RDMA硬件的能力?首先，我們來分析一下深度學(xué)習(xí)應(yīng)用的幾個(gè)特點(diǎn)：
1)Tensor是深度學(xué)習(xí)計(jì)算中最主要的數(shù)據(jù)結(jié)構(gòu)，大量的計(jì)算開銷都是花在對Tensor的處理上。Tensor是一種比較簡單的數(shù)據(jù)結(jié)構(gòu)，主要由meta-data和payload兩部分組成。Payload就是基本元素的數(shù)組，而meta-data就是Tensor的shape信息，也就是維度和每一維的大小。這種簡單的數(shù)據(jù)結(jié)構(gòu)在傳輸?shù)臅r(shí)候其實(shí)不太需要復(fù)雜的序列化和反序列化的功能。
2)在相當(dāng)多的情況下，Tensor是稠密的，并且其大小也是比較大的，也就是說在傳輸這樣的Tensor的時(shí)候并不需要對其進(jìn)行額外的批處理。
3)深度學(xué)習(xí)的訓(xùn)練過程是迭代的。每個(gè)迭代處理一個(gè)mini-batch。在不同的迭代之間，數(shù)據(jù)流圖和很多Tensor的shape信息并不發(fā)生改變，并且其中不少的shape信息是可以在運(yùn)行時(shí)前就靜態(tài)決定的。
基于以上幾個(gè)特點(diǎn)，我們可以對數(shù)據(jù)流圖進(jìn)行分析，找到那些可以靜態(tài)決定shape信息的Tensor，以便在運(yùn)行前，在接收端預(yù)先為其分配RDMA可訪問的內(nèi)存空間，并將其相應(yīng)的可遠(yuǎn)程訪問的地址傳送給發(fā)送端。這樣一來，在運(yùn)行時(shí)，發(fā)送端可以通過單邊的RDMA請求將Tensor的數(shù)據(jù)直接傳輸?shù)浇邮斩�，從而完全避免了沒有必要的額外內(nèi)存拷貝，達(dá)到零拷貝的通信過程。我們將這種機(jī)制在TensorFlow上進(jìn)行實(shí)驗(yàn)， 和基于TCP/IP的gRPC相比，這一方法在一系列典型模型上均取得了多倍的性能改進(jìn)。甚至和針對RDMA優(yōu)化過的gRPC相比，我們的方法仍然能夠取得超過50%的性能提升。
另外，我們在分布式深度學(xué)習(xí)方向上關(guān)注的另一個(gè)問題是如何自動(dòng)地對資源無關(guān)的數(shù)據(jù)流圖做優(yōu)化的分布式執(zhí)行，也就是自動(dòng)劃分?jǐn)?shù)據(jù)流圖中的計(jì)算任務(wù)并為其分配相應(yīng)的計(jì)算資源，以使計(jì)算效率最優(yōu)化。Google的Jeff Dean團(tuán)隊(duì)在這個(gè)方向上已經(jīng)做了很好的先驅(qū)性工作。但局限于模型并行和單機(jī)多卡的運(yùn)行環(huán)境，目前這仍然是一個(gè)非常重要并且大有可為的方向，需要結(jié)合數(shù)據(jù)并行，分布式及異構(gòu)環(huán)境來綜合考慮。
提升單個(gè)計(jì)算單元的運(yùn)算效率
前面提到過，使用深度學(xué)習(xí)框架來實(shí)現(xiàn)的模型算法，在運(yùn)行時(shí)前會(huì)被轉(zhuǎn)換成數(shù)據(jù)流圖。不少具有實(shí)際應(yīng)用價(jià)值的模型都非常復(fù)雜，由它們所轉(zhuǎn)換出來的數(shù)據(jù)流圖通常是由成千上萬的操作節(jié)點(diǎn)構(gòu)成，其中包含了很多運(yùn)算量非常小的節(jié)點(diǎn)，也就是說它們的輸入矩陣的大小很小，或者是其計(jì)算邏輯的復(fù)雜度相對于對輸入數(shù)據(jù)訪問的復(fù)雜度來說很低。大量這樣的操作節(jié)點(diǎn)會(huì)引入以下一些運(yùn)行時(shí)開銷，并且這樣的開銷會(huì)非常顯著。
1)深度學(xué)習(xí)系統(tǒng)運(yùn)行時(shí)需要根據(jù)數(shù)據(jù)流圖中節(jié)點(diǎn)的依賴關(guān)系來調(diào)度節(jié)點(diǎn)的執(zhí)行。調(diào)度每個(gè)節(jié)點(diǎn)的系統(tǒng)開銷和操作節(jié)點(diǎn)計(jì)算量的大小并沒有直接關(guān)系，因此對于由許多小的操作節(jié)點(diǎn)構(gòu)成的計(jì)算流圖來說，系統(tǒng)調(diào)度所帶來的額外開銷就會(huì)相對比較大;
2)對于在GPU上運(yùn)行的計(jì)算來說，每個(gè)操作節(jié)點(diǎn)的實(shí)現(xiàn)都對應(yīng)著一個(gè)GPU的內(nèi)核函數(shù)，而這個(gè)內(nèi)核函數(shù)的每一次執(zhí)行需要CPU調(diào)用顯卡驅(qū)動(dòng)來啟動(dòng)，因此也帶來了常數(shù)量級的額外開銷。這個(gè)開銷相對于計(jì)算量小的內(nèi)核函數(shù)的執(zhí)行來說是非常明顯的;
3)計(jì)算量小的操作節(jié)點(diǎn)往往難以挖掘出足夠的數(shù)據(jù)并行性，因此不能充分利用處理器硬件中的計(jì)算資源。
解決這一問題的主要思路是內(nèi)核融合(Kernel Fusion)。一些手工的優(yōu)化方法就運(yùn)用了這一思想，比如NVIDIA基于CuDNN的RNN庫函數(shù)。它把整個(gè)循環(huán)神經(jīng)網(wǎng)絡(luò)實(shí)現(xiàn)成一個(gè)GPU的內(nèi)核函數(shù)，因此獲得了非常好的性能。然而它的缺點(diǎn)也非常明顯，那就是不夠靈活和通用，無法應(yīng)用在其它網(wǎng)絡(luò)或一些變種的循環(huán)神經(jīng)網(wǎng)絡(luò)中。而我們更加關(guān)注的是如何在深度學(xué)習(xí)的系統(tǒng)中自動(dòng)地對任意的網(wǎng)絡(luò)模型實(shí)施優(yōu)化。
目前在學(xué)術(shù)界和工業(yè)界已經(jīng)存在一些系統(tǒng)采用編譯的方法生成融合的內(nèi)核代碼，比如TVM、Halide和Taco等。這些系統(tǒng)使用Tensor Algebra作為前端表示方法，每個(gè)Tensor Algebra表達(dá)式進(jìn)而可以被編譯成相應(yīng)的內(nèi)核代碼。而Tensor Algebra可以作為更低一層的中間表達(dá)被集成到深度學(xué)習(xí)系統(tǒng)中，也就是說高層的數(shù)據(jù)流圖可以先轉(zhuǎn)換成由Tensor Algebra表達(dá)式組成的代碼塊，再被編譯成可執(zhí)行的代碼。然而，這些系統(tǒng)對于可以進(jìn)行融合的操作節(jié)點(diǎn)有很多限制，不能很好地融合多個(gè)非pointwise的操作，例如多個(gè)矩陣乘操作。然而，我們發(fā)現(xiàn)如果打破這一限制從而融合更多操作節(jié)點(diǎn)是可以帶來更多顯著的性能提升的。
在GPU的運(yùn)行環(huán)境下融合多個(gè)非pointwise的操作具有一定的挑戰(zhàn)性，因?yàn)榉莗ointwise的操作中輸入矩陣的每個(gè)元素都可能依賴于前一個(gè)操作的輸出矩陣中的許多不同位置的元素值，所以在這兩個(gè)操作之間需要插入Barrier同步原語。而在GPU中實(shí)現(xiàn)Barrier需要保證該內(nèi)核的所有線程塊在運(yùn)行時(shí)都是保持活動(dòng)狀態(tài)的，這意味著我們必須要求融合后的內(nèi)核采用有限個(gè)數(shù)的線程塊，但同時(shí)又能夠處理遠(yuǎn)超過線程塊數(shù)量的數(shù)據(jù)塊。
為了解決這一問題，我們嘗試采用persistent-thread的線程塊模型，也就是說在融合后的內(nèi)核的整個(gè)生命周期啟動(dòng)固定數(shù)目的線程塊并讓它們保持活動(dòng)狀態(tài)。我們的優(yōu)化系統(tǒng)在產(chǎn)生融合的內(nèi)核代碼的過程中類似于解決一個(gè)裝箱(bin-pack)問題，即把待融合的子數(shù)據(jù)流圖中的每一個(gè)操作節(jié)點(diǎn)所要處理的數(shù)據(jù)塊分派給適當(dāng)?shù)幕顒?dòng)線程塊，從而使得每個(gè)線程塊的負(fù)載盡可能均衡，并且保持操作節(jié)點(diǎn)的運(yùn)算在原數(shù)據(jù)流圖中的并行性。
為了生成優(yōu)化的GPU內(nèi)核函數(shù)，一個(gè)重要的考慮因素是線程塊和數(shù)據(jù)塊的合理劃分。然而這又依賴于一些非常復(fù)雜的因素，比如操作節(jié)點(diǎn)運(yùn)算中計(jì)算和訪存復(fù)雜度的比率、GPU的shared memory的大小、寄存器文件的大小及分配方法等等。因此一個(gè)最優(yōu)的選擇是很難通過靜態(tài)的方法決定的。幸運(yùn)的是，深度學(xué)習(xí)的迭代性以及需要相當(dāng)多的迭代才能收斂的特性使得我們可以利用早期的迭代過程來收集運(yùn)行時(shí)的動(dòng)態(tài)信息以幫助優(yōu)化系統(tǒng)做更明智的決定。
克服設(shè)備內(nèi)存資源限制
設(shè)備內(nèi)存的大小往往限制了可以處理的模型規(guī)模，解決這一問題的一個(gè)思路是對模型進(jìn)行壓縮和量化。如今學(xué)術(shù)界和工業(yè)界已經(jīng)有大量的研究工作提出不同的壓縮和量化的方法，然而，在實(shí)際的應(yīng)用場景中使用壓縮和量化仍然是個(gè)繁瑣的迭代過程。在這個(gè)過程中，用戶可能會(huì)進(jìn)行以下幾個(gè)方面的嘗試。
1)不同的壓縮方法。比如，是根據(jù)模型的參數(shù)值是否趨近于零，還是將其轉(zhuǎn)換成某種貢獻(xiàn)值之后趨近于零?壓縮時(shí)是不是考慮一定的結(jié)構(gòu)化(如果是面向GPU，可能需要壓縮成塊狀稀疏矩陣來提高運(yùn)行效率)?量化的值點(diǎn)是根據(jù)值域平均劃分還是基于某種聚類來劃分?
2)不同的壓縮程度。要考慮在哪些層的神經(jīng)元參數(shù)上做壓縮，因?yàn)椴⒉皇撬袑訉嚎s后模型效果的敏感程度是一樣的;選擇不同的壓縮率或量化的比特?cái)?shù)。
3)為了保持在大的壓縮率下仍然取得好的模型效果，壓縮過程可能需要是漸進(jìn)的，比如一次壓縮10%，然后重新訓(xùn)練，重復(fù)此過程直到取得目標(biāo)的壓縮率。那么每次漸進(jìn)過程的壓縮率就是一個(gè)需要調(diào)整的參數(shù)。
顯然，這樣一個(gè)繁瑣的過程需要一個(gè)好的工具來使之變得方便。這也是我們組正在關(guān)注的一個(gè)問題。我們正在嘗試擴(kuò)展TensorFlow的API來使用戶可以在模型腳本中直接控制量化和壓縮的方法、對象、程度和過程。
壓縮和量化通常是用來解決模型部署時(shí)的性能和內(nèi)存資源不足的問題，而解決模型訓(xùn)練時(shí)內(nèi)存不夠的問題的思路之一是用計(jì)算來換內(nèi)存。比如，如果數(shù)據(jù)流圖中某一個(gè)操作節(jié)點(diǎn)的計(jì)算量很小，但是輸出的中間結(jié)果數(shù)據(jù)量很大，一個(gè)更好的處理方式是不在內(nèi)存中保存這個(gè)中間結(jié)果，而在后面需要用到它的時(shí)候再重新執(zhí)行這個(gè)操作節(jié)點(diǎn)的計(jì)算。當(dāng)然，重新計(jì)算還是引入了一定的額外開銷。
事實(shí)上，還存在另外一種解決這個(gè)問題的思路，就是將大的輸入數(shù)據(jù)就保存在CPU端的主存里，并將操作節(jié)點(diǎn)實(shí)現(xiàn)成流式的處理，將大的輸入數(shù)據(jù)分段拷貝進(jìn)GPU的設(shè)備內(nèi)存，并通過異步的拷貝使得對每一分段的計(jì)算時(shí)間和下一分段的拷貝時(shí)間能夠重疊起來，從而掩蓋住數(shù)據(jù)拷貝的開銷。對于矩陣乘法這樣的操作，由于計(jì)算復(fù)雜度相對于訪存復(fù)雜度較高，當(dāng)分段較大的時(shí)候，計(jì)算時(shí)間和拷貝時(shí)間是可以達(dá)到完美重疊的。然而，如果所要進(jìn)行的操作不是矩陣乘法，而是一些簡單的pointwise操作，計(jì)算的復(fù)雜度就沒有辦法和內(nèi)存拷貝的開銷相抵消。所以這種做法還需要跟內(nèi)核融合相結(jié)合。比如將矩陣乘法和后續(xù)的pointwise操作相融合，每一個(gè)分段的計(jì)算都會(huì)把該分段的矩陣乘和pointwise操作都做完，然后再處理下一個(gè)分段。
