[翻譯] TensorFlow 分散式之論文篇 "Implementation of Control Flow in TensorFlow"

羅西的思考發表於2022-03-15

分散式

[翻譯] TensorFlow 分散式之論文篇 "Implementation of Control Flow in TensorFlow"

[翻譯] TensorFlow 分散式之論文篇 "Implementation of Control Flow in TensorFlow"

讀論文有一種原則是：本領域最經典的論文，近5年最熱的論文，近1年最新的論文。按照這個原則，本文主要介紹一篇Tensorflow 經典論文 Implementation of Control Flow in TensorFlow。

1. 概覽

本文介紹了 TensorFlow 中控制流操作符的當前設計和實現。這是一份基於原始設計的描述性文件，具體細節請參見實際原始碼。本文內容是：

介紹五個 TensorFlow 的核心操作符，它們是專門為處理控制流而新增的。
展示高層控制流結構如何基於這五個基礎操作符被編譯進資料流圖。
解釋這些資料流圖如何由 TensorFlow runtime 執行，包括在一組混合裝置（如CPU、GPU和TPU）上的分散式執行方式。
描述如何對控制流結構進行自動求導。

本文圖均來自原始論文。

2. 控制流原語

TensorFlow 中控制流的基本設計原則是：引入一個包含少量操作的簡單原子操作集，在這些操作符之上來表達TensorFlow 應用的複雜控制流。我們希望這些基元是靈活且富有表現力的，可以作為高階領域特定語言（DSL）的一個良好的編譯目標。它們應該與 TensorFlow 的資料流模型相相容，並且可以方便實施並行，分散式執行以及自動微分。如下圖所示，原子操作集之中有五個控制流原語運算子，其中 Switch 和 Merge 組合起來可以實現條件控制。所有五個基元一起組合則可以實現 while 迴圈。

圖 1 基元

在 TensorFlow 中，每個 op 都在一個執行幀（execution frame）中執行，控制流原語負責建立和管理這些執行幀。對於每個 while 迴圈，TensorFlow 執行時會設定一個執行幀，並在執行幀內執行 while 迴圈的所有操作。執行幀可以巢狀。巢狀的 while 迴圈在巢狀的執行幀中執行。只要執行幀之間沒有資料依賴關係，則來自不同執行幀的操作可以並行執行。

Switch：Switch 運算子會根據輸入控制張量 p 的布林值，將輸入張量 d 轉發到兩個輸入中的一個。只有兩個輸入都準備好之後，Switch 操作才會執行。

Merge：Merge 運算子將其可用的輸入之一轉發到其輸出。只要它的任何一個輸入可用，merge 運算子就會執行。如果有多個可用的輸入，則無法確定它的輸出。

Enter(name)：Enter 操作符將其輸入轉發到由給定名稱唯一標識的執行幀。這個 Enter 操作用於將一個執行幀中的張量傳遞給一個子執行幀。對於同一個子執行幀可以有多個 Enter 操作，每個操作都會使子執行幀中的張量可用（非同步）。當輸入可用時，Enter 操作將執行。一個新的執行幀在執行該幀第一個 Enter 操作時候被例項化。

Exit：Exit 操作符將一個張量從一個執行幀返回給它的父執行幀。一個執行幀可以有多個 Exit 操作返回到父執行幀，每個操作都非同步地將張量傳回給父幀。當一個 Exit 的輸入可用時，該 Exit 操作就被啟用。

NextIteration: 一個 NextIteration 操作符將其輸入轉發到當前執行幀的下一個迭代。TensorFlow 執行時會跟蹤維護執行幀中的迭代資訊。一個執行幀中執行的任何操作都有一個唯一的迭代 ID，這使得我們能夠唯一地識別迭代計算中同一操作的不同呼叫（比如 hile 操作之中，某一個 op 可能會多次執行）。請注意，一個執行幀中可以有多個 NextIteration操作。當執行幀的第 N 次迭代的第一個 NextIteration 操作開始執行時，TensorFlow 執行時就開始進行第 N+1 次迭代。隨著更多的張量通過執行 NextIteration 操作進入下一個迭代，新迭代中更多操作就開始執行。當一個 NextIteration 的輸入可用時，它就被啟用。

3. 控制流結構的編譯

因為增加了這 5 個控制原語，例如 cond 和 while_loop 這樣的高階程式設計結構就可以被編譯成資料流圖，從而可以被 TensorFlow 執行。我們接下來看看條件表示式和 while 迴圈如何在 Tensorflow 內部實現。

3.1 條件表示式

下面是構建條件表示式 cond(pred, fn1, fn2) 資料流圖的高階虛擬碼。為了簡單起見，我們忽略了實際實現中的許多重細節。讀者可以在 control_flow_ops.py 中找到相關的實現細節。

# Build the graph for the true branch 
context_t = CondContext(pred, branch=1) 
res_t = context_t.Call(fn1)

# Build the graph for the false branch 
context_f = CondContext(pred, branch=0) 
res_f = context_f.Call(fn2)

# Add the Merge nodes for the outputs
merges = [Merge([f, t]) for (f, t) in zip(res_f, res_t)] 
return merges

對於條件表示式的每一個分支，我們都會為條件語境建立一個新的控制流上下文，並在上下文中呼叫其計算圖建構函式（fn1或fn2）。條件上下文允許我們捕獲任何外部張量（不是在上下文中建立的），並插入一個適當的Switch 操作來確保其進入一個分支。這保證了分支中的任何操作只有在該分支被選擇時才會執行。由於 TensorFlow 模型的非同步執行特點，這些外部張量可能在非常不同的時間變得可用，所以我們為每個外部張量使用一個 Switch op 來最大化並行度。

因為每個分支返回一個張量列表（ref_t或res_f），所以我們需要新增一個 Merge 操作來對該結果列表每個輸出的真值/假值進行合併。同樣，輸出可能在不同的時間被計算，所以我們對每個輸出使用一個 Merge 操作，這使我們能夠儘快啟用下游的計算。讓我們來看一個簡單的例子：

圖 2 條件表示式

tf.cond(x<y, lambda: tf.add(x,z), lambda: tf.square(y))

在生成的資料流圖中，Switch 操作被用來控制張量 x、y和z 的流動。在 true/false 分支中，只使用 Switch 操作的真/假輸出。由於 add 的輸入來自 Switch 操作的 true 分支輸出，所以 add 操作只在 x<y 為真時執行。同樣地，Square 操作只在 x<y 為假時執行。Add 或 Square 的結果由最後的 Merge 操作發出。如果條件表示式有多個輸出，就會有多個 Merge 操作，每個輸出都有一個 Merge 操作結果。

有很多種使用 Switch 和 Merge 對 cond 進行編碼的方法，我們選擇目前的編碼方式主要是因為它使 cond 自動求導變得更簡單。

3.2 while 迴圈

以下是構建 while 迴圈資料流圖的高層虛擬碼：

while_context = WhileContext()
while_context.Enter()

# Add the Enter nodes for each loop variable.
enter_vars = [Enter(x, frame_name) for x in loop_vars]

# Add the Merge nodes. Note that input[1] will be updated later.
merge_vars = [Merge([x,x]) for x in enter_vars]

# Build the loop pred subgraph.
pred_result = pred(*merge_vars)

# Add the Switch nodes.
switch_vars = [Switch(x, pred_result) for x in merge_vars]

# Build the loop body subgraph.
body_result = body(*[x[1] for x in switch_vars])

# Add the NextIteration nodes.
next_vars = [NextIteration(x) for x in body_result]

# Form the cycles for the loop.
for m,v in zip(merge_vars, next_vars):
    m.op._update_input(1,v)

# Add the Exit nodes.
exit_vars = [Exit(x[0]) for x in switch_vars]
while_context.Exit()
return exit_vars

整個 while 迴圈圖是在 while 迴圈的控制流上下文之中建立的。這裡的基本思路很簡單。

從迴圈變數開始，我們為每個迴圈變數新增一個 Enter 操作，其後面跟著一個 Merge 操作。然後我們使用其結果（merge_vars）來建立 pred 子圖，pred 子圖將計算迴圈的終止條件。

在加入 Switch 操作後，我們使用 Switch 的 true 分支輸出來構建 while 迴圈主體的子圖。迴圈主體的結果需要進入下一個迭代，所以我們新增 NextIteration 操作，並將其輸出連線到 Merge 操作的第二個輸入。這就形成了迴圈，這使我們在執行圖的時候可以多次重複執行同一個操作。

Switch 操作的假值輸出是整個 while 迴圈的輸出，所以我們在假值輸出後面插入了 Exit 操作，並返回 Exit 操作的輸出。與 cond 類似，while 迴圈的上下文被用來跟蹤 pred 和 body lambdas 中使用的外部張量。這些外部張量被視為迴圈常量，我們為每個這樣的外部張量自動插入一個 Enter 操作，使其可以在 while 迴圈上下文中訪問。巢狀迴圈需要新增巢狀的 Enter 操作。

同樣，讓我們看看一個簡單程式的生成圖例子。

圖 3 while 迴圈

tf.while_loop(lambda i:i<10, lambda i: tf.add(i,1),[0])

在這個例子中，我們只有一個迴圈變數。如果有多個迴圈變數，我們需要新增多個 Enter、Merge、Switch、NextIteration 和 Exit 操作。這樣就可以並行執行跨迴圈和迴圈內跨迭代的操作。我們省略了在 while 迴圈中如何處理常量的方法。如果你想了解其細節，請看具體程式碼。

cond 和 while_loop 的這種轉換方法可以支援條件表示式和迴圈的任意巢狀。例如，一個迴圈體可以呼叫另一個 while_loop，它將被遞迴地翻譯成一個巢狀的子圖。該翻譯確保每個迴圈被靜態地分配一個唯一的框架名稱。

4. 實現

TensorFlow 執行時負責資料流圖的執行。讓我們先快速瀏覽一下。為了在多個裝置上執行，TensorFlow 會自動將操作分配到裝置集上。TensorFlow 基於裝置的具體放置來自動將資料流圖分割成一組子圖，每個裝置一個子圖。當一條邊被分割槽切分時，我們會自動插入一對傳送和接收節點，用於在裝置間傳輸張量。一對 send 和 recv 使用一個唯一的 key 進行通訊，recv 會主動從 send 中提取資料（這裡是特色）。例如，下圖是將一個圖劃分到兩個裝置上的結果，TensorFlow 對分割槽沒有施加任何限制。只要某個節點的計算可以在一個裝置上完成，它就可以被分配到該裝置上。

圖 4 劃分後的計算圖

當一個子圖被分配到某一個裝置之後，這個子圖就被該裝置的本地執行器管理。執行器從源節點開始，依次執行準備好的節點。除了合併節點外，一個節點在其所有輸入都可用時，就成為就緒節點。注意，子圖中的所有 recv 節點都被認為是源節點。

如果沒有控制流，圖的執行就非常直接。每個節點都僅僅被執行一次，當所有節點都被執行過之後，執行就結束了。控制流引入了相當的複雜性。一個節點現在可以被執行任何次數（包括 0 在內）。執行器需要能夠管理同一節點內多個例項的執行（可能是併發的），並確定圖執行何時會完成。

為了跟蹤執行過程中產生的張量，我們使用一個元組 d = (value, is_dead, tag) 來標示執行器中的張量，其中 value 是實際的張量，is_dead 是一個布林值（用來表示該張量是否在一個未執行的條件分支上），而 tag 是唯一標識該張量（以及產生該張量的節點的執行例項）的字串。直觀地說，tag 定義了一個執行環境，在一個執行環境中，一個節點最多執行一次。標籤是傳送/轉發之間通訊 key 的一部分，以區分同一傳送/轉發節點之間的多個呼叫。執行者遵循以下執行規則（注意：一個節點的所有輸入必須有相同的標籤。）

Switch(p,d) = (r1,r2)
r1 = (value(d), p || is_dead(d),tag(d))
r2 = (value(d), !p || is_dead(d),tag(d))

Merge(d1,d2) = r
r = if is_dead(d1) then d2 else d1

Enter(d, frame_name) = r
value(r) = value(d)
is_dead(r) = is_dead(d)
tag(r) = tag(d)/frame_name/0

Exit(d) = r
value(r) = value(d)
is_dead(r) = is_dead(d)
tag(r) = tag1 where tag(d)=tag1/frame_name/n

NextIteration(d) = d1
value(d1) = value(d)
is_dead(d1) = is_dead(d)
tag(d1) = tag1/frame_name/(n+1) where tag(d) = tag1/frame_name/n

Op(d1,...,dm) = (r1,...,rn)
value(ri) = Op.Compute(value(d1),...,value(dm)) if !is_dead(ri)
is_dead(ri) = any(is_dead(d1),...,is_dead(dm)), for all i
tag(ri) = tag(d1), for all i

最後一條規則是針對所有非控制流節點的。請注意，只有當所有的輸入都有效時，才會進行實際的計算。如果有一個無效輸入，我們將跳過計算並向下遊傳播一個 dead 訊號。這種 dead 訊號的傳播可以被用來支援控制流的分散式執行。

5. 分散式條件表示式

對於分散式執行來說，一個條件表示式可能被切分到多個裝置上，如下圖所示：

圖 5 切分表示式

由於任何 recv 節點都是一個隨時無條件啟動的源節點，所以，即使裝置 B 上的 recv 節點是在條件表示式的未選擇分支之內，它也可能會執行。為了使未選擇分支上的 recv 的執行合理化，我們在裝置間把 is_dead 標誌通過 send 節點傳送到 recv 節點。傳播可以在任何數量的裝置上繼續進行。這個簡單的傳播機制可以處理巢狀條件的分散式執行，也有助於 while 迴圈的分散式執行。

6. 分散式的 while 迴圈

對於分散式執行，一個 while 迴圈，特別是迴圈主體，可以被切分到多個裝置上。如果我們簡單地應用切分方案：只是為跨裝置的邊插入 send/recv 節點，那麼裝置上的本地執行器將缺少足夠的資訊來正確執行 while 迴圈。

圖 6 切分控制流簡單方案

讓我們用一個簡單的例子來說明這些問題。在上面的例子中，Op 在迴圈體中，被分配給裝置B。一個簡單切分會將 Switch 到 Op 的邊拆分，插入一對 send/recv 節點，由這對節點完成跨裝置資料傳輸。然而，這是不可行的，因為裝置 B 不知道 recv 和 Op 節點是一個 while 迴圈的一部分，這樣裝置 B 在一個迭代後就會終止執行。解決方案是重寫資料流圖，在每個分割槽新增一個控制迴圈狀態機（如下圖裝置 B 的右下角所示）。控制迴圈 Enter 節點是一個標量 0。

圖 7 切分控制流改進方案

這些控制迴圈提供了足夠的資訊，這樣通過傳送/接收節點相互通訊，就可以使裝置上的執行器能夠像以前一樣獨立執行。請注意，圖中的虛線是控制邊。讓我們先看一下基本用例，即 while 迴圈只執行 0 次迭代。

在裝置 A 上，節點 Enter、Merge、P 和 Switc 依次被執行。因為 P 是 false，所以連線到 Switch 的 Send 會向裝置 B 傳播一個死訊號，這樣 Exit 也會執行，從而使迴圈之外依賴這個 Exit 的節點能夠同時執行。連線到P 的 Send將向裝置 B 傳送布林張量 False，這樣 Recv 也可以被執行，其會等待來自裝置 B 的值。
在裝置 B 上，Enter 觸發了迴圈，接下來依次執行節點 Enter 和 Merge。Merge 的執行使兩個 Recv 得以執行。Switch 的 Recv 會收到 False，所以 Next 會得到一個死張量，於是停止了迴圈。Op 的 Recv 會得到一個死張量，所以 Op 的 Send 會把一個死張量送回裝置 A，此時，裝置 B 沒有未完成的操作，所以執行結束。
在裝置 A 上，Recv for Next 得到了一個死張量。Next 執行，由於它停止了死迴圈的傳播，裝置 A 沒有未完成的操作，所以執行結束。

我們接下來看看 while 迴圈執行一個或多個迭代。

在裝置 A 上，由於 P 在第一次迭代時為真，一個實數張量被髮送到裝置 B。同時 Recv 被執行，等待來自裝置B 返回的值。
在裝置 B 上，控制迴圈狀態機執行並啟用 Recv。Recv 為 Op 從裝置 A 得到一個實數張量；Op 被執行，Send 將一個實數張量送回裝置 A。執行 Next 和 Merge，進一步啟用下一個迭代的 Recv。
在裝置 A 上，Recv 得到一個實數張量。然後執行 Next、Merge 和 P。根據 P 的值，將執行基本情況或新的迭代。

請注意，在執行過程中存在大量的並行性。例如，裝置 B 一旦收到 P 的值，就可以開始下一個迭代或退出。一個參與裝置可以有多個迭代在並行執行，而且兩個參與裝置可以同時在同一個迴圈的不同迭代中工作。

分散式執行 while 迴圈的開銷是每個參與裝置在每次迭代時都需要從產生 P 的裝置那裡接收一個布林張量，考慮到執行中的並行性，開銷在很大程度上應該是與計算重疊，因此可以忽略。

下面顯示了當一個 while 迴圈被劃分到多個裝置上時，資料流圖是什麼樣子的。一個控制迴圈被新增到每個分割槽中，並控制 while 迴圈中的 Recvs。重寫後的圖在語義上與原始圖是等價的。

圖 8 重寫的計算圖

對於巢狀的 while 迴圈，我們按如下方式把控制迴圈堆疊起來。注意，如果一個裝置只有外層迴圈的節點，我們將不會在其上新增任何與內層迴圈有關的控制迴圈結構。

圖 9 巢狀

7. 自動微分

TensorFlow 支援自動求導。例如，使用者可以定義一個帶有損失函式的神經網路，而 TensorFlow 將自動推導並構建反向傳播資料流圖。本節解釋了 TensorFlow 如何在有 cond 和 while_loop 的情況下自動構建反向傳播圖。我們假設讀者對自動反向傳播的工作方式有一定的瞭解。(參見連結 [1]，這是一篇關於反向傳播的優秀文章）。

反向傳播演算法以反向順序遍歷前向圖中的操作，並通過呼叫操作註冊的梯度函式逐步構建梯度圖。一個操作的梯度函式定義了計算該操作梯度的子圖。梯度函式可能會使用到運算的輸入/輸出值，因此在前向計算中產生的一些張量將被保留一段時間，直到它在反向傳播之中被使用。例如，下面顯示了一個前向運算和它的梯度圖。G(Op) 是Op 的梯度子圖。x 和 y 的值將被儲存在記憶體中，直到 G(Op) 被執行。

圖 10 反向傳播

一旦構建了整個資料流圖，TensorFlow 執行時就會自動對圖進行分割，並將執行分佈在多個裝置上。因此，TensorFlow 中的梯度計算也將被分配到多個裝置上執行。

直觀地講，在 cond 和 while_loop 的上下文之中，控制流運算元的反向傳播以如下方式進行反向傳播。Exit 的梯度是 Enter；Switch 的梯度是 Merge（對於cond）或者 NextIteration 之後接著一個 Merge（對於while_loop）；Merge 的梯度是 Switch；NextIteration 的梯度是 Identity；Enter 的梯度是 Exit。TensorFlow 支援巢狀條件和while迴圈的反向傳播。

7.1 條件表示式的反向傳播

直觀地說，cond(p, fn1, fn2) 的梯度為 cond(p, g_fn1, g_fn2)，其中 g_fn1 和 g_fn2 分別為 fn1 和 fn2 的梯度。下面顯示了當 cond 沒有巢狀在 while 迴圈中，cond 的基本反向傳播操作。我們假設 Op 位於 cond 的 true 分支上。如果 cond 被巢狀在 while 迴圈，那麼它需要做更多的工作來記住前向迴圈每次迭代的 p 值。我們將在後面看while 迴圈的反向傳播時討論這個問題。

圖 10 條件表示式的反向傳播

前向傳播之中的 Merge 在後向傳播之中被轉化為 Switch，它使用與前向 Switch 相同的謂詞 p。梯度 g 被反推到Switch 的兩個分支。

前向 Switch 被轉化為 Merge。如果前向 Switch 中只有一個分支在前向傳播之中被用到了，我們會新增一個零輸入到反向傳播的 Merge，如下圖所示，以確保在反向傳播之中總有一個活躍的梯度流經 Merge。這個零輸入被一個 Switch 來控制，所以它只在 p 為 false 時才會被髮送到 Merge。

圖 12 Switch 轉換

7.2 While 迴圈的反向傳播

直觀地說，while_loop(pred, body) 的梯度也是以 while loop 的形式存在。

def pred(i, _): return i < N

while_loop(pred, g_body, [0] + g_vars)

其中 N 是前向傳播 while 迴圈執行的迭代次數，g_body 是前向迴圈體的梯度，g_vars 是迴圈變數的初始值。我們將在後面看到，g_vars 包括前向 while 迴圈變數的初始梯度。下面是一個 while 迴圈的前向傳播和反向傳播圖。

圖 13 While 迴圈的反向傳播

請注意，Backprop 迴圈由 N 控制，即前向迴圈執行的迭代次數。這意味著我們假設 pred 是不可訓練的。G(Body) 是 Body 的梯度。Body 可能再次包含 while 迴圈，所以這個結構可能會遞迴地出現，以處理巢狀的 while 迴圈。

到目前為止，這個描述是相當過度簡化了。實際上，在圖的構造過程中，N 並不是靜態已知的。更重要的是，G(Body) 可能會使用前向傳播過程中產生的值，我們希望保留這些值，以避免在反推過程中重新計算它們。解決方案是重寫前向 while 迴圈的圖，對於反向傳播之中需要的值，增加計算和/或儲存的邏輯。

為了計算 N，我們在前向 while 迴圈中加入以下子圖（計算 N 的邏輯）。因此，N 將由前向迴圈動態計算，並作為後向迴圈的計數迴圈變數的初始值。

圖 14 計算邏輯

為了在反向傳播迴圈中重用前向傳播計算出來的數值，我們在構建反向傳播 while 迴圈的過程中，自動檢測反向傳播中需要的前向值。對於每個這樣的前向值 x，我們自動引入一個堆疊，並在前向迴圈中新增節點，以便在每次迭代時將其值儲存到堆疊中。反向傳播迴圈以相反的順序使用堆疊中的值。堆疊位於前向和反向傳播迴圈之外，由兩個迴圈共享（所以下圖有兩個 Enter）。

圖 15 迴圈共享

實際的計算圖構造實際上比這更微妙和複雜。下面是一些問題。

為了保證正確性，我們需要確保堆疊的 push 和 pop 是按其各自迴圈的迭代來排序的。我們還需要確保前向傳播的堆疊必須在後向傳播的堆疊之前完成排序。這些順序是通過控制邊來完成的。
為了提高效能，我們使堆疊 push 和 pop 操作成為非同步的，因此它們可以與實際計算並行執行。例如，op（甚至是未來的迭代）可以與 push 並行執行。
如果 op 在一個巢狀在 while 迴圈內的 cond 裡面，那麼入棧和出棧操作必須由 cond 的謂詞進行適當的保護。
如果某個值在反向傳播之中被縮減操作（如 Shape、Rank或Size）處理，我們將縮減操作移到前向迴圈中以減少記憶體的使用。

如前所述，Enter 的梯度是 Exit。對於迴圈變數，這就是它的全部作用。對於迴圈常量，我們還新增了一個子圖來累積它們的梯度，如下圖所示。

圖 16 累計梯度

假設 x 是前向傳播中的一個迴圈常數。在 Backprop 中，每次迭代都會為 x 產生一個 partial gradient。因此，我們在反向傳播過程中新增小的累積子圖，然後將所有這些部分梯度加在一起。最終結果 \(g_x\) 是所有偏導數的總和。注意，積累是 eagerly 地進行的，以並行迭代的次數為界。這與 static unrolling 不同，在 static unrolling 中，AddN 需要所有的部分梯度在同一時間生效。

這種結構對巢狀條件和迴圈都有效。對於巢狀在 while 迴圈中的條件式，我們引入一個堆疊來儲存每次前向迭代的謂詞值，並在反向 prop 中使用堆疊中的值（以相反的順序）。對於巢狀的迴圈，當我們遇到巢狀在迴圈體中的內部 while 迴圈時，會遞迴地呼叫這個結構。

一個重要的優化是記憶體交換（memory swapping）。正如我們所看到的，對於每個在 backprop 中需要的前向值 v，我們將其在所有迭代中的值 \(v_1,...,v_N\)儲存在一個堆疊中，所以我們會在 backprop 中重使它們。這對於在記憶體有限的裝置（如GPU）上進行訓練是一個限制。我們使用記憶體交換來非同步地將儲存在堆疊中的值從 GPU 移動到 CPU，並在 Backprop 中需要時將它們移回 GPU 記憶體中。