Jeff Dean等提出動態控制流編程模型，大規模機器學習性能提高21%

時間 2019-12-06

標籤 jeff dean 提出動態控制編程模型大規模機器學習性能提高 21% 欄目系統性能简体版

原文原文鏈接

策劃編輯 | Natalie

編譯 | 蓋磊

編輯 | Natalie

AI 前線導讀：一些最新機器學習模型，特別是 RNN 和強化學習，其訓練和推理中使用了細粒度的動態控制流。在這些模型中，循環體執行和依賴於數據的條件執行須要動態控制流的支持，實如今分佈異構環境各計算設備（包括 CPU、GPU，以及 TPU 等定製的 ASIC）間做快速的控制流決策。Google Brains 團隊認爲，出於性能、可擴展性和運算表達能力上的考慮，現代機器學習系統應該支持分佈式異構計算環境中的動態控制流。

本文是一篇 EuroSys 2018 會議接收論文，論文的主要貢獻者爲 Yuan Yu（本文工做主要完成於在 Google Brain 工做期間），以及 Google Brain 團隊 Jeff Dean 等人，該團隊實現並開源了 TensorFlow 系統。

論文提出了一種支持動態控制流的分佈式機器學習編程模型，並介紹了該模型的設計及在 TensorFlow 中的實現。爲表示機器學習運算，該編程模型對數據流圖（data graph）作了擴展，實現了多個顯著的特性。第一，該編程模型支持對數據流圖的細粒度分區，實現了在多個異構設備上運行條件分支和循環體；第二，使用該模型編寫的程序支持自動微分（automatic differentiation）計算和分佈式梯度計算，這些計算是訓練機器學習模型訓練所必須的；第三，模型採用了非嚴格語義（non-strict semantics），支持多個循環迭代在設備間的並行執行，實現了計算和 I/O 操做的交疊（overlapping）。

Synopsys 的唐杉博士點評道：「傳統神經網絡大部分是靜態圖，這也是不少專用加速器可以高效加速神經網絡的一個基本前提。若是是一個動態的控制流，則底層架構設計可能有巨大差異，運算、存儲和通訊這三要素在硬件中的比例又會發生新的變化。」本文提出的動態控制流編程模型，將給 AI 硬件設計帶來新的挑戰。

更多優質內容請關注微信公衆號「AI 前線」，（ID：ai-front）

背景

機器學習及應用的日新月異，對支持運算的底層系統提出了新的挑戰。人們但願計算系統可擴展運行於不一樣設備上，小到我的的智能手機，大到整個數據中心，並有效地利用各類計算資源，包括 CPU、GPU，以及 TPU 等定製 ASIC。算法

圖 1 給出的是 Google 天然語言處理所用模型的一個簡化架構，其中包括了兩個 RNN 層，以及一個提供 RNN 層間動態學習鏈接的 MoE（混合專家，Mixture of Experts）層。架構中的各個組件可運行在不一樣的設備上，其中使用了基本的動態控制流。計算的不一樣實現策略，對模型的性能具備很大的影響。例如，實現高性能 RNN 的挑戰在於如何權衡內存的使用和計算時間。若是使用多 GPU 的分佈式計算，在 GPU 間交換張量可下降對內存的使用，這爲權衡添加了新的維度。數據庫

圖 1：一個模型架構例子編程

當前，機器學習的構件模塊（例如，上例中的單個單元），以及使用這些模塊所組成的架構，都正在快速地發展中，而且這一趨勢也將繼續。與將 RNN、MoE 等特性層實現爲編程模型原語的方法相比，一種更好的作法是實現通用的控制流構件，例如，分支條件和循環計算。一個機器學習系統，應該提供一種通用的動態控制流機制。緩存

在現代機器學習架構中，一般將計算表示爲數據流圖，並由一個客戶進程驅動一組加速器（accelerator）執行各個構件的計算。所使用的方法能夠分爲兩大類：服務器

圖內實現方式（in-graph）：控制流的決策和操做是編碼實如今數據流圖中的；微信
圖外實現方式（out-of-graph）：由獨立的客戶進程實現控制流決策，而操做則使用 Python 等主機語言提供的控制流原語實現。網絡

在 TensorFlow 的設計中，採用了圖內方式。圖內方式在編譯和運行時間上具備優點。對於運行機器學習模型的集羣而言，模型處理統一的數據流圖，所以能夠優化整個程序。特別是，圖內方式能夠肯定模型運算中間結果的輕重緩急，並據此做出決策，是否須要緩存中間結果。此外，圖內控制方法支持在系統運行時內執行整個計算，這很是適合於異構環境，由於在異構環境中客戶進程間的通訊和同步的代價很是大。若是不使用圖內動態控制流特性，那麼 TensorFlow 程序必須使用 Python 等主機語言所提供的控制流特性。對於循環計算而言，圖內方式比圖外方式可得到更大的並行性。基於本文的實驗結果，對於一臺具備 8 個 GPU 的單機，每秒迭代次數可實現五倍的提高。架構

論文的主要工做可概括爲：併發

論文提出了五種簡單並強大的原語，支持高層控制流構件進一步編譯爲新的原語，實現了更細粒度的數據流圖分區，進而可經過輕量級的協調機制在一組異構設備上執行。app
模型支持自動微分和梯度計算。自動微分是機器學習模型訓練中的重要技術。論文所提出的編程模型，支持在數據流圖中添加計算梯度的子圖，並實現了計算的優化和內存管理技術。計算子圖也可進一步分區運行在多臺設備上。
模型支持並行和異步，提供了非嚴格語義操做，一旦輸入可用時，即調用相應的操做，實現調節分支和循環操做。這種設計支持了 CPU 上的控制流邏輯、GPU 上的計算內核和 CPU 與 GPU 間的內存拷貝間的交疊。

論文提出的編程模型已實如今 TensorFlow 中，並已成功應用於 Google 去年超過 1170 萬的機器學習任務，其中近 65% 的任務包含一個以上的循環。結果顯示，模型的性能和可擴展性適合於真實系統上的真實應用。

編程模型

出於性能和異構運行的須要，TensorFlow 核心運行時使用 C++ 實現。TensorFlow 對多種語言提供了 API，這提供了構建數據流圖的高層接口，並可管理數據流圖在一組設備上的執行。TensorFlow 運行時負責執行數據流圖。爲支持異構系統上的分佈式執行，TensorFlow 具備一箇中央協調器（coordinator），它自動將圖中的節點映射爲一組設備。並將圖中的邊替換爲一對通訊操做，Send(t,k)和Recv(k)。通訊操做對共享同一約會密鑰k，傳遞張量t。數據流圖分區後，每一個子圖將提交給相應的設備，由設備本地運行的執行器執行。各個本地執行器間使用Send(t,k)和Recv(K)通訊，無需中央協調器參與。張量t一旦生成並須要在設備間傳遞，Recv(K)便可從發送設備拉取。

在數據流圖中引入動態控制流，提出了一些額外的挑戰。若是不存在控制流，那麼圖中的每一個操做只會執行一次，每一個值具備惟一命名，而且通訊對使用惟一的約會密鑰。可是若是存在控制流，例如循環，那麼一個操做將可能屢次執行。這時不能作惟一命名和惟一密鑰，它們必須動態生成，以區分對同一操做的屢次調用。

條件分支計算cond(pred,true_fn,false_fn)，其中pred是一個布爾張量，true_fn和false_fn是子圖的構成函數，返回一個張量元組。cond返回一個張量元組，表示執行條件分支計算的結果。它僅使用了Switch、Merge。
循環計算while_loop(pred,body,inits)表示了一個迭代計算，其中pred和body是子圖中循環計算的構成函數，分別是循環終止條件和循環體，輸入參數爲循環變量元組，body返回更新後的循環變量元組。inits是一個張量元組，指定循環變量的初始值。圖 3 是 while 循環的實現。

TensorFlow 的一個特性是不對分區作任何限制。不管圖的拓撲如何，只要設備有能力運行相關操做，就可將操做賦予該設備。分區條件分支和循環體也能夠任意分區，並運行在多個設備上。所以，必須在條件分支計算中實現通知機制，用於告知任何等待Recv的未執行操做從新請求資源。而對於迭代操做，必須實現一種機制，告知參與循環的分區是啓動下一輪計算，或是終止計算。

設計與實現

模型設計中的關鍵點在於：如何用原語表示數據流圖中的分支和循環基本計算，以支持更細粒度的分區；如何用分支和循環表示自動微分和分佈式梯度計算；如何提升模型的處理能力。

控制流原語

在模型的設計上，論文考慮提供一小組靈活的、有表現力的原語，做爲計算數據流模型中高層控制流構件的編譯目標。模型給出了五個控制流原語，分別是Switch、Merge、Enter、NextIteration和Exit，如圖 2 所示。高層控制流構件可編譯爲由上述控制流原語組成的數據流圖，進而實現數據流圖的細粒度分區。

圖 2：控制流原語

使用論文提出的原語，可將while_loop(pred,body,inits)表示爲以下數據流圖。

圖 3：While 循環的數據流圖

當本地執行一個數據流圖時，分區子圖中的全部Recv節點視爲輸入節點。一旦源節點的輸入就緒就執行節點（Merge除外）。若是數據流圖中沒有控制流構件，那麼每一個節點只執行一次。當全部節點完成後，執行完成。若是存在控制流構件，動態控制流引入了新的複雜性。一個節點可能會被屢次執行。執行器必須管理同一操做的屢次執行實例，並根據終止條件斷定整個執行的結束。所以，本地執行器須要從新設計，以處理動態控制流。爲處理同一操做不一樣調用所產生的不一樣張量，執行器內每一個張量表示爲tuple(value,isdead，tag)，其中value是張量的值，is_dead是表示張量是不是Switch未處理分支的布爾值，tag是全局惟一的張量標識符。在 TensorFlow 中，每次循環迭代啓動一個新的幀（frame），在每一個幀中，一個操做最多執行一次。所以，tag能夠區分不一樣迭代所生成的張量。這對於正確的Send和Recv操做十分關鍵，由於tag就是約會密鑰。

當分佈式執行一個數據流圖時，因爲數據流圖被分區爲多個子圖，實現動態控制流的挑戰在於條件分支或循環體子圖被分區到不一樣的設備上。設計中，每一個執行器在執行分區時是獨立相互通訊的，無需中央協調器的參與。在循環執行時，無需設備間的同步，不然將會極大地限制並行性。本地執行器間經過Send和Recv通訊，中央協調器只有在事件完成或失敗時才參與。

對於條件分支，未選取分支上的Recv操做一直處於等待狀態，並阻塞執行，以避免從新申請資源。若是相應的Send操做並未執行，將在設備間從Send到Recv傳播is_dead信號。若是Send-Recv操做對不少，傳播會致使性能開銷。這種狀況儘管不多發生，可是須要對此作一些優化。

對於循環的分佈執行，在每次迭代中，每一個分區須要知道是繼續執行，仍是須要退出。在實現中，模型將自動重寫數據流圖爲簡單的控制狀態機。圖 4 顯示了在兩個設備間分區一個基本的 while_loop 的狀況。循環體包括一個分配給設備 B 的Op操做。這樣，在 B 的分區中，添加了一個循環控制狀態機，用於控制循環體中的Recv操做。圖中的點虛線是控制邊，它肯定了操做的執行順序。循環體分佈式執行的開銷在於，每次迭代中每一個參與設備須要從生成循環條件的設備上接收一個布爾值。因爲通訊是異步的，循環條件的計算一般要領先於其他計算。所以從整個模型看開銷是很小的。例如，對於 GPU 執行器，控制流決策是由本機 GPU 的本地執行器作出的。這樣，本地執行器對於計算和 I/O 操做是徹底併發運行的，並使用各自的資源。所以，動態控制流能夠給出與靜態展開 (static unrolling) 相同的性能（靜態展開是在定義模型建立數據流圖的時候，序列的長度是固定的，以後傳入的全部序列的長度都必須是定義時指定的長度。內存的使用，隨序列長度呈線性增加）。

圖 4：while_loop 的分佈式執行

自動微分的實現

在機器學習算法中，常使用基於梯度的方法優化一組參數。模型訓練期間，近大半的計算時間是用於梯度計算的。所以，提升梯度計算的效率和擴展性很是關鍵。TensorFlow 支持自動微分計算。給定一個表示了神經網絡計算的數據流圖，TensorFlow 將生成實現分佈式梯度計算的高效代碼。下面介紹如何將自動微分計算表示爲控制流構件。

TensorFlow 提供了一個反向自動微分庫（autodiff），實現了後向傳播算法。下面介紹如何支持cond和while_loop構件。tf.gradients()函數計算標量函數 f(x1,x2,...) 的張量，其中 x1,x2,... 是一組張量參數，算法實現了向量的鏈式法則（chain rule）。圖 5 給出了梯度函數G(Op)的數據流圖結構，函數依賴於 y 的偏微分，偏微分基於原始Op函數輸出 gz 及輸入，即以張量表示的矩陣 x 和 y。因爲 x 和 y 要被梯度函數G(Op)使用，所以在操做執行期間須要一直保持在內存中。由此所致使的內存佔用，是影響深度神經網絡訓練能力的一個關鍵因素。

圖 5：操做 Op 及其梯度函數 G(Op)

那麼動態控制流構件是如何支持 TensorFlow 的 autodiff 算法執行後向傳播的？數據流圖中，每一個操做關聯到一個「控制流上下文」，它指定了最內層的控制流構件，其中的操做是由各個構件組成。當後向傳播變量首次趕上的新控制流上下文時，它就在梯度計算數據流圖中生成一個相應的控制流構件。這樣，輸出梯度g_z的操做tf.cond(pred,true_fn,false_fn)，其梯度爲tf.cond(pred,true_fn_grad(g_z),false_fn_grad(g_z))。而對於tf.while_loop(cond,body,loop_vars)，下面的代碼顯示瞭如何靜態展開循環體，並對tf.matmul()和tf.reduce_sum()應用梯度計算函數。

能夠看到，前向循環的中間結果a_1、a_2和a_3將在梯度計算循環中使用。計算的性能取決於如何處理這些中間結果。爲此，TensorFlow 引入了一種新的棧數據結果，用於保存循環間的中間值。前向計算將結果壓入棧中，而梯度計算則從棧中彈出值。圖 6 給了實現這種基於棧的狀態保存的數據流圖結構。

圖 6：在後向傳播中保存可重用的張量

內存管理

對於 GPU 等特定設備，內存管理在上述過程當中十分關鍵。由於 GPU 的內存使用一般侷限於 16GB 之內。本論文在模型的實現中採用了多種技術。其中一種主要的技術使用了臨時本地性，稱爲「內存交換」（memory swap）。當一個張量壓入堆棧時，內存交換技術將其從 GPU 移動到 CPU 內存中，並在後向傳播使用該張量時逐步放回。內存交換的高效實現，關鍵在於計算和 I/O 操做間的交疊。這須要多個系統組件的無縫合做，包括：

循環中的屢次迭代能夠並行運行，堆棧壓入和彈出操做是同步的，可在計算中並行運行。
GPU 對計算和 I/O 操做使用用獨立的 GPU 流，以增進這兩類操做的交疊。
內存交換僅在內存使用超出設定閾值時啓用。

實驗代表，若是不使用內存交換機制，那麼系統在處理序列長度 500 時就無內存可用。但在啓用了內存交換後，同一系統能夠處理序列長度 1000，而且開銷很小。可用的內存狀況，是限制系統處理最大序列長度能力的首要因素。

評估與實際應用

實驗運行使用了一套生產集羣，該集羣是與其它任務共享的，由 Intel 服務器組成，配置了 Nvidia Tesla K40 GPU，經過以太網鏈接。論文中評估了模型實現的幾個關鍵特性，包括性能、可擴展性、內存交換特性、序列長度處理能力等。

基準測試：模型性能和可擴展性

圖 8 和圖 9 是對分佈式迭代計算性能和擴展能力的基準測試狀況。基準測試採用了一個 while_loop，循環體分區運行在一個 GPU 集羣上。如圖 7 所示，測試中採用的一種數據同步狀況是設備間不作協調，另外一種是每次迭代結束後全部設備等待協調（相似於 AllReduce）。圖 8 左圖顯示了機器數從 1 增長到 64，每秒實現迭代次數的變化狀況。圖 8 右圖代表，並行運行迭代對於實現 GPU 間的並行十分關鍵。圖 9 給出了 GPU 數量對於處理性能（表示爲每秒迭代次數）的影響狀況。

圖 7：分佈式 while-loop 中數據流依賴性

圖 8：分佈式 while-loop 的性能分析（左圖更改機器數量，右圖更改迭代次數）

圖 9：計算性能（以每秒迭代次數表示）隨 GPU 數量的變化狀況。

模型特性測試

表 1 是使用內存交換並增長訓練長度的運行狀況。

表 1：LSTM 模型每次循環迭代訓練時間，增長序列長度

圖 10 對比了使用動態控制流和靜態展開的狀況。

圖 10：動態控制流與靜態展開的性能對比

實際應用：使用動態控制流實現強化學習

爲展現動態控制流的優勢，論文給出了一個強化學習實現的實例。做爲基準測試，文中實現了深度 Q 網絡（DQN，Deep Q-Networks），動態控制流如圖 11 所示。數據庫用於存儲輸入的經驗，並經過週期性地對數據庫採樣，提供以往的經驗給 Q 學習算法。Q 學習算法使用這些經驗，構建另外一個稱爲「目標網絡」的神經網絡去訓練主神經網絡。目標網絡也是週期性更新的，以反映主網絡的週期性快照狀況。動態控制流支持整個計算保持在系統運行時中，並支持並行計算。相比於基準系統，動態控制可提升性能 21%。同時，用戶無需使用主機語言對數據流圖各個分區作編程，簡化了算法的編程實現。