[翻譯] TensorFlow 分散式之論文篇 "TensorFlow : Large-Scale Machine Learning on Heterogeneous Distributed Systems"

羅西的思考 發表於 2022-03-12
Machine Learning

[翻譯] TensorFlow 分散式之論文篇 "TensorFlow : Large-Scale Machine Learning on Heterogeneous Distributed Systems"

本系列我們開始分析 TensorFlow 的分散式。之前在機器學習分散式這一系列分析之中,我們大多是以 PyTorch 為例,結合其他框架/庫來穿插完成。但是缺少了 TensorFlow 就會覺得整個世界(系列)都是不完美的,不單單因為 TensorFlow 本身的影響力,更因為 TensorFlow 分散式有自己的鮮明特色,對於技術愛好者來說是一個巨大寶藏。

讀論文有一種原則是:本領域最經典的論文,近5年最熱的論文,近1年最新的論文。按照這個原則,本文主要介紹一篇 TensorFlow 經典論文 TensorFlow : Large-Scale Machine Learning on Heterogeneous Distributed Systems。大家如果讀了下面論文就會發現 TensorFlow分散式的博大精深。

本文圖來自原始論文。

1. 原文摘要

基於我們在 DistBelief 方面的經驗,以及對期望的系統特性和訓練/使用神經網路需求的更全面的理解,我們構建了 TensorFlow ,這是我們的第二代系統,用於實現和部署大規模機器學習模型。 TensorFlow 採用類似資料流的模型來描述計算,並將其對映到各種不同的硬體平臺上,從在 Android 和 iOS 等移動裝置平臺上執行推理,到使用包含一個或多個 GPU 卡的單機的中等規模訓練和推理系統,再到在數百臺具有數千個 GPU 的專用主機上執行的大規模訓練系統。

擁有一個能夠跨越如此廣泛的平臺的單一系統可以顯著簡化機器學習系統的實際使用,因為我們發現,如果大規模訓練和小規模部署都分別有自己的獨立系統,則會導致巨大的維護負擔和漏洞。 TensorFlow 計算被表示為有狀態資料流圖,我們致力於使系統具有足夠的靈活性,以便使用者可以快速試驗新模型,系統同時也具有足夠高的效能和魯棒性,可以被用於機器學習模型的訓練和部署。

為了將神經網路訓練擴充套件到更大規模的部署, TensorFlow 允許客戶機通過複製和並行執行核心模型資料流圖來輕鬆表達各種並行性,這樣可以使用許多不同的計算裝置來更新一組共享引數或其他共享狀態。對計算描述的適度更改允許使用者實現多種不同的並行方法。

TensorFlow 允許在引數更新的一致性方面具有一定的靈活性,這些寬鬆的同步要求允許我們可以在一些較大的部署中更加輕鬆。與 DistBelief 相比, TensorFlow 的程式設計模型更加靈活,效能顯著提高,支援在更廣泛的異構硬體平臺上訓練和使用更廣泛的模型。

2. 程式設計模型和基本概念

TensorFlow 計算由一組節點組成的有向來描述。該圖表示一個資料流計算,也允許讓某些型別的節點來維護和更新持久狀態,並以類似於 Naiad 的方式在圖中實現分支和迴圈控制結構。客戶端通常使用一種前端語言(C++或Python)構建計算圖。下圖顯示了使用 Python 前端構建並執行 TensorFlow 計算圖的示例片段,

[翻譯] TensorFlow 分散式之論文篇 "TensorFlow : Large-Scale Machine Learning on Heterogeneous Distributed Systems"

圖 1. TensorFlow 計算圖示例片段

[翻譯] TensorFlow 分散式之論文篇 "TensorFlow : Large-Scale Machine Learning on Heterogeneous Distributed Systems"

圖 2,計算圖

在 TensorFlow 圖中,每個節點表示操作的例項,其具有零個或多個輸入和零個或多個輸出。在計算圖中沿普通邊流動的值(從輸出到輸入)被稱為張量。張量是任意維陣列,其基本元素型別在計算圖構造時被指定或推斷出來。特殊的邊(被稱為控制依賴關係)也可以存在於圖中:這些邊上沒有資料流,但它們表明控制依賴關係的源節點必須在控制依賴關係的目標節點開始執行之前完成執行。因為我們的模型包含可變狀態,因此客戶可以直接使用控制依賴關係來強制執行。我們的實現有時也會在內部插入控制依賴項,以強制某些獨立操作之間的順序,例如,可以控制峰值記憶體使用。

2.1 運算元(Operations)與核(Kernels)

運算元(Operation)表示一個抽象計算(例如,"矩陣乘法"或"加法")。一個運算元可以擁有屬性,但是所有屬性必須在計算圖構造時被提供或推斷出來,這樣才能例項化一個執行該運算元的節點。屬性的一個常見用途是使運算元在不同的張量元素型別上具有多型性(例如,加法運算元即支援兩個型別為 float 的 tensors 相加,也支援兩個型別為 int32的張量相加)。

核(Kernel)是可以在特定型別的裝置(例如CPU或GPU)上執行的運算元的具體實現。 TensorFlow 通過序號產生器制定義了一系列運算元和核,這樣意味著使用者可以通過連結其他運算元和/或核心來進行擴充套件。下圖顯示了 TensorFlow 庫中內建的一些運算元。

[翻譯] TensorFlow 分散式之論文篇 "TensorFlow : Large-Scale Machine Learning on Heterogeneous Distributed Systems"

表 1. 運算元

2.2 會話(Sessions)

客戶端程式通過建立會話與 TensorFlow 系統互動。會話先構建一個空計算圖,為了建立計算圖,會話介面支援Extend 方法,該方法可以讓客戶端附加節點和邊來擴充這個空圖,然後進行計算。

2.3 變數(Variables)

在大多數計算中,一個圖會被執行多次,而大多數張量在圖的單次執行之後都不會存在。然而,有一些張量是在計算圖的執行過程之中始終存在的,位置也是固定的,其不能正常流動但是可以更新,比如模型的引數,這就引出了變數這個概念。

變數是一種特殊的操作,它返回持久可變張量的控制程式碼,這些控制程式碼可以被傳遞給少量特殊的操作,例如 AssignAssignAdd(相當於+=),通過這些操作就可以改變這些變數引用的張量。

3. 實現

TensorFlow 系統中的主要元件是客戶端,它使用會話介面與主機以及一個或多個工作程式進行通訊。每個工作程式負責協調對一個或多個計算裝置(如 CPU 核心或 GPU 卡)的訪問以及按照主裝置的指示在這些裝置上執行計算圖節點。 TensorFlow 介面有本地分散式實現兩種。

  • 當客戶端、master 和 worker 都在單個機器上單個程式的上下文之中執行時(如果機器安裝了多個 GPU 卡,則可能使用多個裝置),將使用本地實現。

  • 分散式實現與本地實現共享大部分程式碼,但支援通過一個環境對其進行擴充套件,在該環境中,客戶端、master和 worker 都可以在不同機器上不同的程式中。在我們的分散式環境中,這些不同的任務是 job 之中的一些容器,這些 job 由叢集排程系統來管理。這兩種不同的模式如下圖所示。

[翻譯] TensorFlow 分散式之論文篇 "TensorFlow : Large-Scale Machine Learning on Heterogeneous Distributed Systems"

圖 3. 排程模式

3.1 裝置(Devices)

裝置是 TensorFlow 的計算核心。每個工作者負責一個或多個裝置,每個裝置都有一個裝置型別和名稱。裝置名稱由以下幾部分組成:

  • 裝置型別。
  • 裝置在工作者中的索引。
  • 分散式設定中對於工作者所在作業和任務的標識(如果裝置是程式本地的,則為 localhost)。

示例裝置名稱的樣例如下:

  • "/job:localhost/device:cpu:0"。
  • "/job:worker/task:17/device:gpu:3"。

PyTorch 有針對 CPU 和 GPU 的裝置介面的實現,其他裝置型別可以通過序號產生器制提供新裝置實現。每個裝置物件負責管理裝置記憶體的分配和釋放,以及執行 TensorFlow 下發的核方法。

3.2 張量

在我們的實現中,張量是一個型別化的多維陣列。我們支援多種張量元素型別,包括大小從 8 位到 64 位的有符號和無符號整數、IEEE 浮點和雙精度型別、複數型別和字串型別(arbitrary byte array)。張量所在裝置的分配器負責管理張量的儲存區,張量儲存緩衝區是引用計數的,在沒有引用保留時會進行釋放。

3.3 單裝置執行

讓我們首先考慮最簡單的執行場景:一個擁有單個裝置的工作者程式。計算圖中的節點按照節點之間依賴關係的順序來執行。我們將跟蹤每個節點尚未執行的依賴項數量的計數。一旦此計數降至零,該節點就有資格執行,並被新增到就緒佇列中。就緒佇列以某種未指定的順序進行處理,其將節點的核方法執行委託給裝置物件。當節點完成執行時,依賴於此已完成節點的所有節點的計數都將減少。

3.4 多裝置執行

一旦一個系統有多個裝置,就有兩個主要的複雜問題:如何決定將每個節點的計算放在哪個裝置上,如何管理這些放置(Placement )所帶來的跨裝置資料通訊。本小節討論這兩個問題。

3.4.1 決定裝置(Node Placement)

給定計算圖之後, TensorFlow 實現的主要職責之一是將計算對映到可用裝置集。本文給出了該演算法的一個簡化版本。

佈局(placement)演算法的一個輸入是成本模型(cost model),該模型包含每個計算圖節點的輸入和輸出張量的大小(位元組)估計,以及每個節點在獲得輸入張量之後所需的計算時間估計。該成本模型要麼基於與不同操作型別相關的啟發式靜態估計,要麼基於計算圖早期執行的實際佈局決策集進行測量/決定。

佈局演算法首先執行計算圖的模擬執行,然後使用貪婪啟發式為圖中的每個節點選擇一個裝置。此模擬生成的"節點到裝置放置關係"最終也用作實際執行的放置。佈局演算法從計算圖的源開始,並在前進過程中模擬系統中每個裝置上的活動,在此遍歷中:

  • 當到達了一個節點,就考慮此節點的可使用裝置集(如果裝置不提供使用者希望的實現特定操作的核心,則裝置就不使用)。
  • 對於具有多個可用裝置的節點,佈局演算法使用貪婪啟發式演算法,看看將節點放置在每個可能裝置上對節點完成時間會造成怎樣的影響。該啟發式方法不僅考慮了在成本模型中,該類裝置上操作的估計或測量執行時間,還包括為了將輸入從其他裝置傳輸到該節點而引入的通訊成本。然後選擇最快完成節點操作的裝置作為該操作的裝置,
  • 佈局(placement)過程繼續進行,以便為圖中的其他節點(包括現在準備好模擬執行的下游節點)做出放置決策。

3.4.2 跨裝置通訊(Cross-Device Communication)

一旦決定了節點如何放置到裝置之上(node placement),圖就被劃分成一組子圖,每個裝置一個子圖。任何跨裝置的從 x 到 y 的邊將被刪除,並替換為兩條新邊:

  • 一條邊是在 x 子圖中,從 x 到新的 Send 節點。
  • 一條邊是在 y 的子圖之中,從對應的 Receive 節點到 y。

請參見下圖。

[翻譯] TensorFlow 分散式之論文篇 "TensorFlow : Large-Scale Machine Learning on Heterogeneous Distributed Systems"

圖 4 插入傳送/接收節點之前和之後

在執行時,傳送和接收節點將會彼此協調如何在裝置之間傳輸資料。這使我們能夠把傳送和接收的所有通訊隔離出來,從而簡化執行時(runtime)的其餘部分。

當我們插入傳送和接收節點時,我們規範如下:特定裝置上特定張量的所有使用者都使用同一個接收節點,而不是特定裝置上的每個下游使用者都擁有一個自己的接收節點。這確保了所需張量的資料在"源裝置→ 目標裝置對"之間只傳輸一次,並且目標裝置上張量的記憶體只分配一次,而不是多次(例如,參見上圖中的節點 b 和 c)。

通過以這種方式處理通訊,我們還允許將不同裝置上的圖的各個節點的排程分散到工作者之中:傳送和接收節點在不同的工作者和裝置之間傳遞必要的同步,這樣就把主節點從排程任務之中解放出來。主節點只需要向每個具有計算圖的任何節點的工作者發出單個 Run 請求(每次計算圖執行),而不需要參與每個節點或每個跨裝置通訊的排程。這使得系統更具可伸縮性可擴充套件性,並且和主節點強制執行排程相比,可以執行更細粒度的節點執行策略。

3.5 分散式執行

計算圖的分散式執行與多裝置執行非常相似。在決定裝置如何放置之後,將為每個裝置建立一個子圖。傳送/接收節點對在跨工作程式通訊時候使用遠端通訊機制(如 TCP 或 RDMA)來跨機器邊界移動資料。

3.5.1 容錯(Fault Tolerance)

我們可以在許多地方檢測到分散式執行中的故障,目前主要依靠如下兩種方式:

  • 檢測到傳送和接收節點對之間的通訊錯誤。
  • 從主程式到每個工作程式的定期健康檢查。

當檢測到故障時,整個計算圖執行將中止並從頭開始。因為變數(Variable)節點指的是在圖的執行過程中持續存在的張量,所以我們支援設定一致性檢查點,以此來在重新啟動時恢復這些狀態。具體來說,每個變數節點都連線到一個 Save 節點。這些儲存節點定期執行,例如每 N 次迭代執行一次,或每 N 秒執行一次。當它們執行時,變數的內容被寫入持久儲存,例如分散式檔案系統。類似地,每個變數也都連線到一個恢復節點,該節點僅在重新啟動後的第一次迭代中啟用。

4. 擴充套件(Extensions)

在本節中,我們將介紹基本程式設計模型的幾個更高階的特性。

4.1 計算梯度

許多優化演算法,包括常見的機器學習訓練演算法(如隨機梯度下降),會使用一組輸入來計算一個成本函式(cost function)的梯度。因為這是一種常見的需求,所以 TensorFlow 內建了對自動梯度計算的支援。如果一個 TensorFlow 計算圖中的張量 C 可能通過一個複雜的操作子圖依賴於一組張量{\(X_k\)},那麼一個內建函式將返回張量集{\(dC/dX_k\)}。梯度張量和其他張量一樣,通過使用以下步驟擴充套件 TensorFlow 圖來計算。

張量 C 依賴於張量 I,當 TensorFlow 需要計算張量 C 相對於張量I的梯度時,它首先在計算圖中找到從 I 到 C 的路徑。然後它從 C 回溯到 I,對於反向路徑上的每個操作,它會向 TensorFlow 圖新增一個節點,使用鏈式規則沿向後路徑合成偏導數。

新新增的節點為前向路徑中的相應操作計算"梯度函式"。梯度函式可以通過任何操作註冊。該函式不僅將沿反向路徑計算的部分梯度作為輸入,還可以選擇正向操作的輸入和輸出。圖5顯示了根據圖2示例計算的成本梯度。灰色箭頭顯示梯度函式的潛在輸入,該函式不用於所示的特定操作。

[翻譯] TensorFlow 分散式之論文篇 "TensorFlow : Large-Scale Machine Learning on Heterogeneous Distributed Systems"

圖 5 計算這些梯度所需的附加值為:

\[[db,dW,dx] = tf.gradients(C,[b,W,x]) \]

自動梯度計算會使優化更加複雜化,尤其是記憶體使用。當執行前向計運算元圖(即使用者顯式構造的子圖)時,啟發式優化演算法可以通過觀察計算圖的構造順序來決定下一步執行哪個節點,這通常意味著臨時輸出在構建後很快就會被佔用,因此它們的記憶體可以很快被重用。

當啟發式無效時,使用者可以通過更改計算圖構造的順序,或新增控制依賴項來優化記憶體使用。但是,當梯度節點自動新增到計算圖中時,使用者的控制能力會降低,啟發式演算法可能會崩潰。特別是,因為梯度反轉了正向計算順序,因此在計算圖執行中,早期使用的張量在梯度計算的末尾經常再次需要。這種張量會佔用大量稀缺的 GPU 記憶體,從而不必要地限制計算量。我們正在積極改進記憶體管理,以便更好地處理此類情況。選項包括使用更復雜的啟發演算法來確定計算圖執行的順序,重新計算張量而不是將其保留在記憶體中,以及將長期張量從 GPU 記憶體交換到更大的主機 CPU 記憶體。

4.2 區域性執行(Partial Execution)

客戶機通常只想執行整個執行圖的一個子圖。為了支援這一點,一旦客戶機在會話中設定了計算圖,我們的 Run 方法允許客戶機執行整個圖的任意子圖,並沿圖中的任意邊輸入任意資料,以及沿圖中任意邊獲取資料。

圖中的每個節點都有一個名稱,節點的每個輸出由源節點名稱和節點的輸出埠標識,從 0 開始編號(例如,"bar:0" 表示"bar"節點的第一個輸出,而"bar:1"表示第二個輸出)。

Run 呼叫的兩個引數有助於定義將要執行的計算圖的確切子圖。首先,Run 呼叫接受型別為 name:port 名稱到"fed"張量值的對映作為輸入。其次,Run 呼叫接受 output names,這是一個輸出 name[:port] 列表,其指定了應執行哪些節點。如果名稱中存在埠部分,則如果 Run 呼叫成功完成,應將節點的特定輸出張量值返回給客戶端。

計算圖可以基於輸入和輸出的值進行轉換。輸入中每個 node:port 都替換為一個 feed 節點,該節點將從用於 Run 呼叫的 Rendezvous 物件中獲取輸入張量。類似地,每個帶有埠的輸出名稱都連線到一個特殊的 fetch 節點,該節點被用來儲存輸出張量,並在執行呼叫完成時將其返回給客戶端。最後,一旦通過插入這些特殊的 feed 和 fetch 節點重寫了計算圖,要執行的節點集可以通過以下方式確定:從每個由輸出指定的節點開始,使用圖依賴關係在圖中進行後向傳播,以確定為了計算輸出而必須在重寫圖中執行的完整節點集。下圖在左側顯示了一個原始圖,以及在使用 inputs{b} 和 outputs{f:0} 呼叫 Run 時生成的轉換圖。因為我們只需要計算節點f的輸出,所以我們不會執行節點 d 和 e,因為它們與 f 的輸出沒有任何關係。

[翻譯] TensorFlow 分散式之論文篇 "TensorFlow : Large-Scale Machine Learning on Heterogeneous Distributed Systems"

圖 6 區域性執行前後

4.3 裝置約束(Device Constraints)

TensorFlow 客戶端可以通過為節點提供部分約束來控制節點在裝置上的位置,這些約束與節點可以在哪些裝置上執行有關。例如,"僅將此節點放置在 GPU 型別的裝置上",或"此節點可以放置在/job:worker/task:17"中的任何裝置上,或"將此節點與名為variable13"的節點合併"。在這些約束條件的約束範圍內,佈局演算法(placement algorithm)負責完成節點到裝置的分配,以提供快速的計算執行,並滿足裝置自身施加的各種約束,例如,限制裝置上執行其計算圖節點子集所需的記憶體總量。

支援此類約束要求更改前面描述的佈局演算法。我們首先計算每個節點的可行裝置集,然後在共定位約束圖(graph of colocation constraints)上使用 union find 來計算出必須放置在一起的圖元件。對於每個這樣的元件,我們計算可行裝置集的交集。

4.4 控制流

雖然沒有任何顯式控制流的資料流圖也非常有表達能力,但我們發現,在很多情況下,如果支援條件和迴圈,則可以用更簡潔和有效來表示機器學習演算法。

與 Arvind 描述的資料流機(dataflow-machine)方法一樣,我們在 TensorFlow 中引入了一小部分控制流原語操作符,並將 TensorFlow 推廣到可以處理迴圈資料流圖。SwitchMerge 運算子允許我們根據布林張量的值來跳過整個子圖的執行。Enter,LeaveNextIteration 運算子允許我們表示迭代。高層級的程式設計結構,如 if-conditionals 和 while-loops 則可以使用這些控制流操作符來輕鬆地編譯成資料流計算圖。

TensorFlow 執行時實現了標籤(tags)和幀(frames)的概念,其在概念上類似於 MIT 標記令牌機(MIT Tagged-Token machine)。迴圈的每個迭代都由一個 tag 唯一標識,其執行狀態由一個 frame 表示。只要輸入準備好,它就可以進入迭代,因此可以同時執行多個迭代。

如何為分散式系統處理迴圈控制的狀態? TensorFlow 使用分散式協調機制來執行帶有控制流的圖。通常,迴圈可以包含分配給許多不同裝置的節點。因此,管理迴圈的狀態成為分散式終止檢測問題。TensorFlow 的解決方案是基於圖重寫。在圖分割槽過程中,我們自動向每個分割槽新增控制節點。這些節點實現了一個小型狀態機,它協調每個迭代的開始和結束,並決定最終迴圈的結束。對於每個迭代,擁有迴圈終止斷言(predicate)的裝置向每個參與的裝置傳送一條控制訊息。

如上所述,我們通常通過梯度下降來訓練機器學習模型,並將梯度計算表示為資料流圖的一部分。當模型包含控制流操作時,我們必須在相應的梯度計算中考慮它們。例如,對於具有 if-conditional 的模型,梯度計算需要知道採用了條件的哪個分支,然後將梯度邏輯應用於該分支。類似地,帶有 while-loop 的模型的梯度計算需要知道進行了多少次迭代,並且還將依賴於在這些迭代期間計算的中間值。目前依然採用重寫計算圖技術來記錄梯度計算所需的值。

4.5 輸入操作

雖然可以通過 feed 節點把輸入資料提供給計算呼叫,但用於訓練大規模機器學習模型的另一種常見機制是在圖中部署有特定的輸入操作節點,這種節點通常配置成一組檔名,該節點每次執行時產生一個張量,該張量包含儲存在該組檔案中的資料的一個或多個樣本。這樣就允許將資料直接從底層儲存系統讀取到計算機記憶體中,然後計算機記憶體將對資料執行後續處理。在客戶端程式與工作程式分開的配置中,如果資料被饋送,則通常需要額外的網路躍點 hop(從儲存系統到客戶端,然後從客戶端到工作程式,而不是使用輸入節點時直接從儲存系統傳輸到工作程式)。

4.6 佇列

佇列是我們新增到 TensorFlow 中的一個有用特性。它們允許計算圖的不同部分進行非同步操作,並通過入隊(Enqueue)和出隊(Dequeue)操作傳遞資料。入隊操作可以阻塞,直到佇列中有可用的空間,而出隊操作也可以阻塞,直到佇列中有所需的最少數量的元素可用。佇列的一種用途是,當機器學習模型的計算部分仍在處理前一批資料時,模型可以從磁碟檔案中預取輸入資料。它們也可用於其他型別的分組操作,包括累積多個梯度,這樣可以把小 batch 組合成為一個大 batch,以便在大的批次上計算更復雜的梯度組合,或將迴圈語言模型的不同輸入句子分組到大致相同長度的句子箱(bin)中,然後可以更有效地處理這些問題。

除了普通的 FIFO 佇列之外,我們還實現了一個 shuffling 佇列,它在一個大的記憶體緩衝區內隨機 shuffle 其元素。例如,對於希望用隨機順序來處理樣本的機器學習演算法來說,這種洗牌功能非常有用。

4.7 容器

容器是 TensorFlow 中用於管理長期可變狀態的機制。預設容器將會一直持續到程式終止,但我們也允許使用其他的命名容器。容器儲存變數的備份,可以通過完全清除容器中的內容來重置容器。通過使用容器可以在不同會話的完全不相交的計算圖之間共享狀態。

5 優化

在本節中,我們將介紹 TensorFlow 實現中的一些優化,這些優化可以提高系統的效能或資源利用率。

5.1 消除公共子表示式

由於計算圖的構造通常由客戶機程式碼中的許多不同層的抽象來完成,因此計算圖很容易存在重複計算。為了處理這個問題,我們實現了類似於Click(原始論文參考文獻12)描述的演算法,該演算法在計算圖上執行,並將具有相同輸入和操作型別的多個操作副本縮減到其中的一個節點,並會把相應的邊進行重定向。

5.2 控制資料傳輸和記憶體使用

仔細安排 TensorFlow 操作可以提高系統的效能,特別是在資料傳輸和記憶體使用方面。具體而言,排程可以減少中間結果儲存在記憶體中的時間,從而減少記憶體消耗峰值。這種減少對於記憶體不足的 GPU 裝置尤為重要。此外,重新安排裝置間的資料通訊可以減少對網路資源的爭奪。

雖然有很多機會可以安排優化,但這裡我們重點關注一個我們認為特別必要和有效的機會。它涉及接收節點讀取遠端值的計劃。如果不採取預防措施,這些節點可能會比必要時啟動得更快,可能在執行開始時一次啟動。通過執行運籌學中常見的儘快/儘可能晚(ASAP/ALAP)計算,我們分析圖的關鍵路徑以估計何時啟動接收節點。然後,我們插入控制邊,目的是在需要這些節點的結果之前取消這些節點的起始位置。

5.3 非同步核

除了正常同步核心之外,我們的框架還支援非阻塞核心。這樣的非阻塞核心使用一個稍有不同的介面,通過該介面,一個 continuation 被傳遞給計算方法,該 continuation 應該在核心執行完成時呼叫。這是一種針對具有多個活動執行緒環境的優化。在這些環境中,記憶體使用或其他資源相對昂貴,因此我們應該避免在等待 I/O 或其他事件發生時長期佔用執行執行緒。非同步核心的示例包括接收核心、排隊核心和出列核心(如果佇列空間不可用或沒有可讀取的資料,則可能需要分別進行阻塞)。

5.4 優化函式庫

目前有很多高度優化的函式庫,比如 BLAS 和 cuBLAS,另外,cuda convnet 和 cuDNN 可以在 GPU 之上處理卷積核。因此我們利用這些函式庫實現了系統的很多核心。

5.5 有失真壓縮(Lossy Compression)

一些機器學習演算法(比如通常用於訓練神經網路的演算法)具有抗噪聲和降低演算法精度的能力。在裝置之間傳送資料時(有時在同一臺機器內的裝置之間),我們通常使用高精度的有失真壓縮,這一方式類似於分散式系統。例如,我們通常插入特殊的轉換節點,將 32 位浮點表示轉換為 16 位浮點表示(不是 IEEE 16 位浮點標準,而是 32 位 IEEE 794 浮點格式,但尾數中的精度降低了 16 位),然後在通訊通道的另一端轉換回 32 位表示(只需為尾數的丟失部分填入零,因為這樣在 32→ 16→ 32 位轉換時計算花費更少)。

6. 程式設計技巧

上面講的都是一些系統級別的優化,下面我們講述一些機器學習演算法所用到的技巧。這裡假定使用者都用 SGD 來求解機器學習演算法。TensorFlow 的基本資料流圖模型可以以多種方式用於機器學習應用。我們關心的一個領域是如何加速計算密集型神經網路模型在大型資料集上的訓練。本節描述了我們和其他人為了實現這一點而開發的幾種技術,並說明了如何使用 TensorFlow 實現這些不同的方法。

本小節中的方法假設使用隨機梯度下降法(SGD)對模型進行訓練,使用的小批次包括 100~1000 個樣本。

6.1 資料並行訓練

可以通過資料並行的方式來提升 SGD 的效率,比如,假如每次 SGD 的 mini-batch 是 1000 個樣本,我們可以把模型複製成 10 份放到 10 個 GPU 之上,mini-batch 也被切成 10 份,每份就是 100 個樣本。這個 10 個模型副本平行計算自己 100 個樣本的梯度,然後把這些梯度進行同步規約,最後用於模型引數更新。這就像我們使用 1000個元素的批量大小執行順序 SGD 演算法一樣。在這種情況下, TensorFlow 圖擁有原計算圖中執行大多數計算部分的多個副本,由單個客戶端執行緒驅動這個大型圖的整個訓練迴圈。下圖頂部對此進行了說明。

[翻譯] TensorFlow 分散式之論文篇 "TensorFlow : Large-Scale Machine Learning on Heterogeneous Distributed Systems"

圖 7 資料並行

這種方法也可以是非同步的,每一個模型副本非同步地將梯度更新應用於模型引數,不需要彼此等待。在此配置中,每個圖副本都有一個客戶端執行緒。上圖下半部分對此進行了說明。

7. Model Parallel Training

模型並行訓練也很容易用 TensorFlow 表示,這樣對於同一批樣本,模型不同部分可以在不同的計算裝置上同時計算。下圖顯示了 LSTM 模型的示例,該模型在三個不同的裝置上並行。

[翻譯] TensorFlow 分散式之論文篇 "TensorFlow : Large-Scale Machine Learning on Heterogeneous Distributed Systems"

0x08 Concurrent Steps for Model Computation PipeLine

在同一裝置中對模型計算進行流水線處理也是一個常用的提高利用率的方法,這是通過在同一組裝置中執行少量的併發步驟來完成的。它有點類似於非同步資料並行,只是流水線並行發生在同一裝置內,而不是在不同裝置上覆制計算圖。在一個單一的步驟中,在所有裝置上的計算可能無法在任何時候完全利用全部裝置的並行性,而流水線並行允許 "填補間隙",這可以充分利用空閒的裝置資源。

[翻譯] TensorFlow 分散式之論文篇 "TensorFlow : Large-Scale Machine Learning on Heterogeneous Distributed Systems"

0xFF 參考

TensorFlow 架構

TensorFlow : Large-Scale Machine Learning on Heterogeneous Distributed Systems

TensorFlow : a system for large-scale machine learning

TensorFlow 分散式採坑記