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

來自專欄 AI前線

本文由 「AI前線」原創（ID：ai-front），原文鏈接：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 間交換張量可降低對內存的使用，這為權衡添加了新的維度。

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

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

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

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

論文的主要工作可歸納為：

• 論文提出了五種簡單並強大的原語，支持高層控制流構件進一步編譯為新的原語，實現了更細粒度的數據流圖分區，進而可通過輕量級的協調機制在一組異構設備上執行。
• 模型支持自動微分和梯度計算。自動微分是機器學習模型訓練中的重要技術。論文所提出的編程模型，支持在數據流圖中添加計算梯度的子圖，並實現了計算的優化和內存管理技術。計運算元圖也可進一步分區運行在多台設備上。
• 模型支持並行和非同步，提供了非嚴格語義操作，一旦輸入可用時，即調用相應的操作，實現調節分支和循環操作。這種設計支持了 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 所示。高層控制流構件可編譯為由上述控制流原語組成的數據流圖，進而實現數據流圖的細粒度分區。

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

當本地執行一個數據流圖時，分區子圖中的所有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) 相同的性能（靜態展開是在定義模型創建數據流圖的時候，序列的長度是固定的，之後傳入的所有序列的長度都必須是定義時指定的長度。內存的使用，隨序列長度呈線性增長）。

自動微分的實現

在機器學習演算法中，常使用基於梯度的方法優化一組參數。模型訓練期間，近大半的計算時間是用於梯度計算的。因此，提高梯度計算的效率和擴展性非常關鍵。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)使用，因此在操作執行期間需要一直保持在內存中。由此所導致的內存佔用，是影響深度神經網路訓練能力的一個關鍵因素。

那麼動態控制流構件是如何支持 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 給了實現這種基於棧的狀態保存的數據流圖結構。

內存管理

對於 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 數量對於處理性能（表示為每秒迭代次數）的影響情況。

模型特性測試

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

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

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

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

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

圖 11：深度 Q-Network 中的動態控制流

結論

論文介紹了一種使用動態控制流實現機器學習的編程模型。該模型設計支持異構設備間的並發和分布執行，可以使用 CPU、GPU 及 TPU 等定製 ASIC 計算，支持應用在多設備間的頻繁控制流決策。模型已在 TensorFlow 中實現，並得到了廣泛地使用，對機器學習的進展做出了貢獻。

動態控制流是機器學習中的一個相當活躍研究領域，對未來工作發展非常重要。尤其是條件計算和流計算，可相應地提高計算能力。近期研究給出的模型具有超過 1000 億個參數。條件計算和流計算的進一步抽象和實現技術值得做進一步研究。

查看論文原文：

https://arxiv.org/pdf/1805.01772.pdf

今日薦文

萬人上書亞馬遜，反對向警方出售人臉識別技術

活動推薦

如果你對人工智慧感興趣，推薦關注《AI 技術內參》。用一年時間，為你精講人工智慧國際頂級學術會議核心論文，系統剖析人工智慧核心技術，解讀技術發展前沿與最新研究成果，分享數據科學家以及數據科學團隊的養成秘笈。

新註冊用戶，立減 30 元。長按下圖二維碼試讀專欄。

向17W+AI愛好者、開發者和科學家，每周一節免費AI公開課，囊括上萬人的AI學習社群，提供最新AI領域技術資訊、一線業界實踐案例、搜羅整理業界技術分享乾貨、最新AI論文解讀。回復「AI前線」、「TF」等關鍵詞可獲取乾貨資料文檔。

如果你希望看到更多類似的優質內容，記得點個贊再走！┏(＾0＾)┛

推薦閱讀：