手把手教你從系統層面優化深度學習計算

來源： 微軟研究院AI頭條

本文約5643字，建議閱讀10分鍾。

在影象、語音識別、自然語言處理、強化學習等許多技術領域中，深度學習是非常有效的，並且某些問題已經達到甚至超越了人類的水平。除了改變模型和演算法，是否可以從系統的層面來優化深度學習計算，進而改善計算資源的使用效率？本文中資深研究員伍鳴與大家分享他對深度學習計算優化的一些看法。

深度學習在近幾年裡取得了巨大的進步，它已經或者是有望成功地被應用在我們許多生活場景中，比如自動駕駛、安防、翻譯、醫療等等。可以說，計算機的計算和通訊能力的大幅提升是促使深度學習成功的重要因素。

深度學習為什麼依賴於超大的計算能力？

首先，深度學習本質上是基於統計的科學，所以大規模的樣本資料對於深度學習的效果是至關重要的。其次，更大規模和更復雜的神經網路模型已經被證明非常有效，並在產品中有廣泛的使用，這同時也產生了對計算能力的更大要求和消耗。舉個例子，具有8層神經元的AlexNet網路2012年在ImageNet資料集上取得16%的錯誤率，該網路的一次迭代執行大約需要1.4 GFLOP的計算量。而微軟提出的使用152層神經元的殘差網路（ResNet）於2015年在該資料集上取得3.5%的錯誤率，其一次迭代的計算量大約是22.6GFLOP，是AlexNet的16倍。在當今的生產環境中，影象、語音以及自然語言處理相關的模型，例如人臉識別、語音轉文字、機器翻譯等，即使給予相當多的計算資源，很多仍需要幾周的時間才能完成訓練。

再次，深度學習模型是迅速迭代的。在AI領域，每年學術界和工業界都會提出大量的新模型。對每一個實際的問題，開發者需要不斷嘗試不同的模型和演算法，甚至對於同一種模型演算法，也需要去反覆除錯超引數以獲得最好的預測效果。可想而知，如果模型的每次訓練都要幾周的時間，那麼尋找最優模型的過程會非常漫長和痛苦。

另外，模型的線上推理具有更加極致的效能要求。線上的服務具有硬性的服務等級協議（SLA），所以在實際部署大型模型時，需要手工重新優化在深度學習框架（如TensorFlow）上已經訓練好的模型，導致大量額外工程開銷的產生。

由此可見，進一步優化深度學習計算對於深度學習的快速發展和成功應用起著至關重要的作用。

深度學習計算優化的挑戰和機會

目前，優化深度學習的計算存在以下幾個主要的挑戰：

• 單機單計算單元（如GPU）的資源限制往往不能滿足對大規模資料和模型的處理要求，那麼就需要使用多機多計算單元來橫向擴充套件計算的規模。如何才能最大限度地減少通訊的開銷從而最大化多機的並行度？

• 如何優化神經網路的計算使得它能夠把單個硬體計算單元的效率發揮到極致？

• 雖然許多硬體計算單元（GPU、FPGA等）的計算能力很強大，但是它們的記憶體資源（即裝置記憶體）非常稀缺。當它們不能提供模型執行所需要的記憶體資源時，要麼運算不能夠進行下去，要麼就需要將計算所需的資料在主存和裝置記憶體之間倒來倒去，帶來很大的執行開銷。如何才能更好地利用有限的裝置記憶體資源從而不給計算效率帶來負面的影響？

• 深度學習開發者和研究人員通常只想關注神經網路模型和演算法本身，並不想被複雜的優化問題分散精力。這意味著深度學習框架這樣的系統軟體最好能夠實現自動優化，而對模型開發者透明。那麼，如何對特定的優化做合理的抽象使其更加靈活通用、更加容易地整合在系統框架中便是需要認真考慮的問題。

事實上，任何方面的優化問題都可以從模型演算法和系統兩個角度來看待。一方面，我們可以通過改變模型和演算法來優化其對計算資源的使用效率從而改進其執行速度。這樣的優化對特定的演算法往往非常有效，但卻不容易擴充套件應用到其它演算法中。而另一方面，也就是微軟亞洲研究院異構計算組正在進行的研究，則是在系統中實施模型演算法無關的優化，這樣的優化，通常可以為更多的應用帶來效能的好處，同時也符合我們在前文提到的透明性的要求。

以系統優化助力深度學習計算

為了能夠更好地理解系統這一層面的優化，我們先來簡單介紹一下深度學習框架系統的背景知識。當今工業界流行的深度學習系統（包括TensorFlow、PyTorch、CNTK、MxNet、Caffe等）大都採用分層的體系結構設計。在前端提供高階語言（例如Python）的介面抽象，允許使用者方便地描述神經網路結構，也就是深度學習的模型。描述好的模型在被系統執行前，首先會被轉換成資料流圖（Data-flow Graph）。在這個資料流圖中，節點是特定的矩陣操作（也就是Operator，如Sigmoid、Matrix Multiplication等），而連線不同節點的邊則是操作節點的輸入和輸出矩陣。這個資料流圖也可以被看成是深度學習計算的中間表達。然後，深度學習系統的後端將這個資料流圖對映到實際硬體上進行高效地執行，而大部分系統層面的優化就是在這個階段完成的。

1. 加速分散式深度學習訓練

分散式訓練的主要瓶頸在於多機之間的通訊開銷。如今計算機網路的硬體技術已經有了很大的發展，InfiniBand的RDMA網路卡（Remote Direct Memory Access，這是一種硬體的網路技術，它使得計算機訪問遠端的記憶體時無需遠端機器上CPU的干預）已經可以提供50~100Gbps的網路頻寬和微秒級的傳輸延遲。目前許多以深度學習為目標應用的GPU機群都部署了這樣的網路。然而深度學習系統如何才能充分利用好硬體提供的通訊能力使分散式的訓練獲得更大的效能提升呢？另外，使用RDMA的軟體介面進行通訊能夠繞過TCP/IP協議棧，減少了作業系統核心態的執行開銷。在這樣的網路通訊技術的支援下，任何與通訊相關的計算處理的開銷都會變得非常顯著，而這正是許多原先基於TCP/IP而設計的網路通訊機制中所存在的問題。

RPC（Remote Procedure Call，遠端過程呼叫）是一個被廣泛使用的多機之間的通訊抽象原語，它的主要設計目標是通用性。在沒有考慮RDMA的情況下，很多深度學習框架都會採用RPC的機制（例如gRPC）來實現多機之間的通訊。然而，RPC需要維護一個內部的私有快取，從而不得不引入使用者資料的儲存空間和內部快取之間的資料拷貝。這種記憶體拷貝的開銷在使用RDMA網路的情況下會變得非常明顯。我們通過micro-benchmark觀察到，跟使用基於TCP/IP的gRPC相比，直接通過RDMA的介面傳輸訊息（對不同的訊息大小）可以有2到10倍的效能提升。

那麼針對深度學習的應用負載，如何才能更好地利用RDMA硬體的能力？首先，我們來分析一下深度學習應用的幾個特點：

• Tensor是深度學習計算中最主要的資料結構，大量的計算開銷都是花在對Tensor的處理上。Tensor是一種比較簡單的資料結構，主要由meta-data和payload兩部分組成。Payload就是基本元素的陣列，而meta-data就是Tensor的shape資訊，也就是維度和每一維的大小。這種簡單的資料結構在傳輸的時候其實不太需要複雜的序列化和反序列化的功能。

• 在相當多的情況下，Tensor是稠密的，並且其大小也是比較大的，也就是說在傳輸這樣的Tensor的時候並不需要對其進行額外的批處理。

• 深度學習的訓練過程是迭代的。每個迭代處理一個mini-batch。在不同的迭代之間，資料流圖和很多Tensor的shape資訊並不發生改變，並且其中不少的shape資訊是可以在執行時前就靜態決定的。

基於以上幾個特點，我們可以對資料流圖進行分析，找到那些可以靜態決定shape資訊的Tensor，以便在執行前，在接收端預先為其分配RDMA可訪問的記憶體空間，並將其相應的可遠端訪問的地址傳送給傳送端。這樣一來，在執行時，傳送端可以通過單邊的RDMA請求將Tensor的資料直接傳輸到接收端，從而完全避免了沒有必要的額外記憶體拷貝，達到零拷貝的通訊過程。我們將這種機制在TensorFlow上進行實驗， 和基於TCP/IP的gRPC相比，這一方法在一系列典型模型上均取得了多倍的效能改進。甚至和針對RDMA優化過的gRPC相比，我們的方法仍然能夠取得超過50%的效能提升。

另外，我們在分散式深度學習方向上關注的另一個問題是如何自動地對資源無關的資料流圖做優化的分散式執行，也就是自動劃分資料流圖中的計算任務併為其分配相應的計算資源，以使計算效率最優化。Google的Jeff Dean團隊在這個方向上已經做了很好的先驅性工作。但侷限於模型並行和單機多卡的執行環境，目前這仍然是一個非常重要並且大有可為的方向，需要結合資料並行，分散式及異構環境來綜合考慮。

2. 提升單個計算單元的運算效率

前面提到過，使用深度學習框架來實現的模型演算法，在執行時前會被轉換成資料流圖。不少具有實際應用價值的模型都非常複雜，由它們所轉換出來的資料流圖通常是由成千上萬的操作節點構成，其中包含了很多運算量非常小的節點，也就是說它們的輸入矩陣的大小很小，或者是其計算邏輯的複雜度相對於對輸入資料訪問的複雜度來說很低。大量這樣的操作節點會引入以下一些執行時開銷，並且這樣的開銷會非常顯著。

• 深度學習系統執行時需要根據資料流圖中節點的依賴關係來排程節點的執行。排程每個節點的系統開銷和操作節點計算量的大小並沒有直接關係，因此對於由許多小的操作節點構成的計算流圖來說，系統排程所帶來的額外開銷就會相對比較大；

• 對於在GPU上執行的計算來說，每個操作節點的實現都對應著一個GPU的核心函式，而這個核心函式的每一次執行需要CPU呼叫顯示卡驅動來啟動，因此也帶來了常數量級的額外開銷。這個開銷相對於計算量小的核心函式的執行來說是非常明顯的；

• 計算量小的操作節點往往難以挖掘出足夠的資料並行性，因此不能充分利用處理器硬體中的計算資源。

解決這一問題的主要思路是核心融合（Kernel Fusion）。一些手工的優化方法就運用了這一思想，比如NVIDIA基於CuDNN的RNN庫函式。它把整個迴圈神經網路實現成一個GPU的核心函式，因此獲得了非常好的效能。然而它的缺點也非常明顯，那就是不夠靈活和通用，無法應用在其它網路或一些變種的迴圈神經網路中。而我們更加關注的是如何在深度學習的系統中自動地對任意的網路模型實施優化。

目前在學術界和工業界已經存在一些系統採用編譯的方法生成融合的核心程式碼，比如TVM、Halide和Taco等。這些系統使用Tensor Algebra作為前端表示方法，每個Tensor Algebra表示式進而可以被編譯成相應的核心程式碼。而Tensor Algebra可以作為更低一層的中間表達被整合到深度學習系統中，也就是說高層的資料流圖可以先轉換成由Tensor Algebra表示式組成的程式碼塊，再被編譯成可執行的程式碼。然而，這些系統對於可以進行融合的操作節點有很多限制，不能很好地融合多個非pointwise的操作，例如多個矩陣乘操作。然而，我們發現如果打破這一限制從而融合更多操作節點是可以帶來更多顯著的效能提升的。

在GPU的執行環境下融合多個非pointwise的操作具有一定的挑戰性，因為非pointwise的操作中輸入矩陣的每個元素都可能依賴於前一個操作的輸出矩陣中的許多不同位置的元素值，所以在這兩個操作之間需要插入Barrier同步原語。而在GPU中實現Barrier需要保證該核心的所有執行緒塊在執行時都是保持活動狀態的，這意味著我們必須要求融合後的核心採用有限個數的執行緒塊，但同時又能夠處理遠超過執行緒塊數量的資料塊。

為了解決這一問題，我們嘗試採用persistent-thread的執行緒塊模型，也就是說在融合後的核心的整個生命週期啟動固定數目的執行緒塊並讓它們保持活動狀態。我們的優化系統在產生融合的核心程式碼的過程中類似於解決一個裝箱（bin-pack）問題，即把待融合的子資料流圖中的每一個操作節點所要處理的資料塊分派給適當的活動執行緒塊，從而使得每個執行緒塊的負載儘可能均衡，並且保持操作節點的運算在原資料流圖中的並行性。

為了生成優化的GPU核心函式，一個重要的考慮因素是執行緒塊和資料塊的合理劃分。然而這又依賴於一些非常複雜的因素，比如操作節點運算中計算和訪存複雜度的比率、GPU的shared memory的大小、暫存器檔案的大小及分配方法等等。因此一個最優的選擇是很難通過靜態的方法決定的。幸運的是，深度學習的迭代性以及需要相當多的迭代才能收斂的特性使得我們可以利用早期的迭代過程來收集執行時的動態資訊以幫助優化系統做更明智的決定。

3. 克服裝置記憶體資源限制

裝置記憶體的大小往往限制了可以處理的模型規模，解決這一問題的一個思路是對模型進行壓縮和量化。如今學術界和工業界已經有大量的研究工作提出不同的壓縮和量化的方法，然而，在實際的應用場景中使用壓縮和量化仍然是個繁瑣的迭代過程。在這個過程中，使用者可能會進行以下幾個方面的嘗試。

• 不同的壓縮方法。比如，是根據模型的引數值是否趨近於零，還是將其轉換成某種貢獻值之後趨近於零？壓縮時是不是考慮一定的結構化（如果是面向GPU，可能需要壓縮成塊狀稀疏矩陣來提高執行效率）？量化的值點是根據值域平均劃分還是基於某種聚類來劃分？

• 不同的壓縮程度。要考慮在哪些層的神經元引數上做壓縮，因為並不是所有層對壓縮後模型效果的敏感程度是一樣的；選擇不同的壓縮率或量化的位元數。

• 為了保持在大的壓縮率下仍然取得好的模型效果，壓縮過程可能需要是漸進的，比如一次壓縮10%，然後重新訓練，重複此過程直到取得目標的壓縮率。那麼每次漸進過程的壓縮率就是一個需要調整的引數。

顯然，這樣一個繁瑣的過程需要一個好的工具來使之變得方便。這也是我們組正在關注的一個問題。我們正在嘗試擴充套件TensorFlow的API來使使用者可以在模型指令碼中直接控制量化和壓縮的方法、物件、程度和過程。

壓縮和量化通常是用來解決模型部署時的效能和記憶體資源不足的問題，而解決模型訓練時記憶體不夠的問題的思路之一是用計算來換記憶體。比如，如果資料流圖中某一個操作節點的計算量很小，但是輸出的中間結果資料量很大，一個更好的處理方式是不在記憶體中儲存這個中間結果，而在後面需要用到它的時候再重新執行這個操作節點的計算。當然，重新計算還是引入了一定的額外開銷。

事實上，還存在另外一種解決這個問題的思路，就是將大的輸入資料就儲存在CPU端的主存裡，並將操作節點實現成流式的處理，將大的輸入資料分段拷貝進GPU的裝置記憶體，並通過非同步的拷貝使得對每一分段的計算時間和下一分段的拷貝時間能夠重疊起來，從而掩蓋住資料拷貝的開銷。對於矩陣乘法這樣的操作，由於計算複雜度相對於訪存複雜度較高，當分段較大的時候，計算時間和拷貝時間是可以達到完美重疊的。然而，如果所要進行的操作不是矩陣乘法，而是一些簡單的pointwise操作，計算的複雜度就沒有辦法和記憶體拷貝的開銷相抵消。所以這種做法還需要跟核心融合相結合。比如將矩陣乘法和後續的pointwise操作相融合，每一個分段的計算都會把該分段的矩陣乘和pointwise操作都做完，然後再處理下一個分段。

作者簡介

伍鳴，微軟亞洲研究院資深研究員。2007年於中科院計算所取得計算機系統結構博士學位後加入微軟亞洲研究院。期間主要的研究興趣及參與的研究方向包括分散式事務處理系統、圖計算引擎和人工智慧平臺。近年來在多個系統領域的頂級會議（如SOSP、OSDI、NSDI、ATC、EuroSys、SoCC、VLDB等）中發表多篇論文，並擔任過OSDI、ASPLOS、HotDep、MiddleWare等會議的程式委員會委員，以及SOSP’17的Publication Chair。