[源碼解析] PyTorch 流水線并行實作 (5)--計算依賴

前幾篇文章我們介紹了 PyTorch 流水線并行的基本知識，自動平衡機制和切分資料等，本文我們結合論文内容來看看如何實作流水線依賴，核心就是如何建立這些小批次之間的跨裝置依賴關系。

0x00 摘要

0x01 前文回顧

0x02 計算依賴

0x03 反向傳播依賴

2.1 解析

2.2 基礎功能

2.2.1 Function

2.2.2 Fork

2.2.3 Join

2.2.4 Phony

2.2.5 detach

2.3 使用

0x03 正向傳播依賴

3.1 分割模型

3.2 建立依賴

0x04 總結

0x06 更正和補充

6.1 緣起

6.2 論文内容

6.2.1 裝置級執行順序

6.2.2 并行計算與拷貝

6.2.3 推論

6.3 實作

6.3.1 _copy_streams

6.3.2 Copy 算子

6.3.3 Wait 算子

6.3.4 封裝

6.3.5 建立依賴關系

6.4 對比圖

6.5 總結

0xFF 參考

流水線并行其他文章連結如下:

[源碼解析] 深度學習流水線并行Gpipe(1)---流水線基本實作

[源碼解析] 深度學習流水線并行GPipe (2) ----- 梯度累積

[源碼解析] 深度學習流水線并行 GPipe(3) ----重計算

[源碼解析] 深度學習流水線并行之PipeDream(1)--- Profile階段

[源碼解析] 深度學習流水線并行 PipeDream(2)--- 計算分區

[源碼解析] 深度學習流水線并行 PipeDream(3)--- 轉換模型

[源碼解析] 深度學習流水線并行 PipeDream(4)--- 運作時引擎

[源碼解析] 深度學習流水線并行 PipeDream(5)--- 通信子產品

[源碼解析] 深度學習流水線并行 PipeDream(6)--- 1F1B政策

[源碼解析] PyTorch 流水線并行實作 (1)--基礎知識

[源碼解析] PyTorch 流水線并行實作 (2)--如何劃分模型

[源碼解析] PyTorch 流水線并行實作 (3)--切分資料和運作時系統

[源碼解析] PyTorch 流水線并行實作 (4)--前向計算

本文圖檔來自論文和github源碼。

為了更好的了解本文，我們首先看看前文之中的關鍵部分。

原始流水線狀态如下：

管道并行的政策是根據分區索引 j 配置設定任務，以便第 j 個分區完全位于第 j 個裝置中。

持有模型後期部分的裝置必須等待，直到持有模型早期部分的裝置計算結束。

目标流水線狀态如下：

目前問題：

如果分成若幹個微批次，則需要強制要求 $F_{i,j}$ 必須在 $F_{i+1,j}$ 之前完成，以及 $B{i,j}$ 必須在執行$B{i-1,j}$ 之前完成。

後向傳播的計算圖是在前向傳播過程中動态構造的。PyTorch既不記錄正向計算圖，也不維護一個梯度錄音帶（gradient tape），PyTorch的自動微分（autograd）引擎僅對計算圖進行反向傳播。這意味着自動加載引擎可能不會完全按照與正向過程相反的執行順序運作，除非由圖的結構強制執行。

目前難點：

如何在每個裝置中以正确的順序釋出那些綁定到裝置的任務，以避免由于Python解釋器未能提前請求而延遲在裝置上（與CPU異步）執行任務。[這個前文已經介紹]

如何建立這些小批次之間的跨裝置依賴關系。

實作方案：

如何保證正确執行順序？torchgpipe引入了确定性時鐘周期（deterministic clock-cycle），它給出了任務的總體順序。[這個前文已經介紹]

如何保證計算圖中的動态顯式依賴關系？針對clock_cycles産生的每一個運作計劃：

利用 fence 函數調用“fork”和“join”，以此在向後計算圖中動态建立顯式後向傳播依賴關系。

利用 compute(schedule, skip_trackers, in_queues, out_queues) 進行計算。

因為前文已經介紹了執行順序方案，是以本文介紹如何計算依賴。

為什麼需要計算依賴？

因為模型已經被分層，模型的不同部分拆開放到不同裝置上，資料也分成微批次，是以本來模型内部是線性依賴關系，現在需要變成流水線依賴關系。是以原始計算圖不能滿足需求，是以需要有針對性的補充。就像上圖那樣，6個層被分成了三個partitions，這三個partitons 之間的依賴如何建構？

之前的線性依賴關系其實是在模型定義時候就基本确定了，現在則需要每次運作時候建立一個動态依賴關系。

是以針對流水線并行，torchgpipe需要自己補充一個本機跨裝置僞分布式依賴關系。torchgpipe 通過在前向計算圖和後向計算圖做各種調整來達到目的。計算圖就意味着各種依賴邏輯，依賴邏輯的補足就是依靠本節介紹的 Fork 和 Join 兩個函數完成的。

這裡最初有一個疑問，就是Torchgpipe怎麼在不使用 PyTorch RPC 和 p2p的情況下，建構出來一個異地反向計算圖。後來發現，原來是我想多了，因為Torchgpipe沒有考慮到這種情況，它針對都是在同一個主機之上的GPU，不涉及異地多機器計算。

Torchgpipe 本質上還是一個程序内運作多個線程進行計算，是 DP 的替代。比如源碼中就有對比如下：

再比如代碼中明确提到：

我們首先看看反向傳播依賴，這個是論文的重點。

我們還是要回憶一下前面兩個圖例。

圖1

圖2

這裡需要完成兩種依賴：

行間依賴，就是 batch 之間的依賴，就是裝置内的依賴。從圖上看，就是藍色列内的 $F_{1,1}$ 必須在 $F_{2,1}$之前完成，$B_{2,1}$ 必須在$B_{1,1}$ 之前完成。

列間依賴，就是 partitions（裝置）之間的依賴。從圖上看，就是藍色 $F_{1,1}$ 必須在黃色 $F_{1,2}$之前完成，即第一個裝置必須在第二個裝置之前完成，而且第一個裝置的輸出是第二個裝置的輸入。

假定我們依據确定性時鐘周期（deterministic clock-cycle）算法來運作一個前向傳播。即使前向傳播是按照在第j個裝置上應該執行的順序來執行任務 $F_{1,j},...,F_{m,j}$ ，得到的後向傳播結果計算圖看起來也更像圖1而非圖2，

從圖1上看，PyTorch 的 autograd 引擎不知道 $B_{i+1,j}$ 必須在 $B_{i,j}$ 之前運作，是以會打亂後向傳播的時間流。是以，虛拟依賴（圖2的虛線箭頭）必須在前向傳播中被顯式繪制出來。

我們再仔細分析一下圖2。圖2之中，每一行都表示一個 micro-batch 在訓練中的運作流，這個流的前向是由clock算法确定的。後向關系是由前向傳播中自動确定完成的。

現在的問題是：一個 mini-batch 被分成了4個 micro-batch，分别在不同時鐘周期進入訓練。就是每一列。這一列由上到下的傳播也是由clock算法确定，但是反向傳播（由下自上）目前是不确定的。比如最後一列中，反向傳播的順序應是：$B_{4,1},B_{3,1},B_{2,1},B_{1,1}$。但是這個目前從前向傳播的結果來看，無法确定這個順序。

是以需要依靠本節介紹的 Fork 和 Join 兩個函數完成這個依賴關系。圖中斜線表示checkpoint之中需要先有一個重計算，然後才能由下往上走。

是以，torchpipe定義兩個基礎函數，Fork 和 Join 來表達這種依賴關系：

Fork 是 auto grad 函數，其把一個張量 x 映射到 pair(x, $\phi$)，這裡 $\phi$ 是一個空張量。

Join 是 auto grad 函數，其把 pair(x, $\phi$) 映射到一個張量 x ，這裡 $\phi$ 是一個空張量。

現在，$F_{i+1,j}$ 對于 $F_{i,j}$ 的依賴（其在後向傳播計算圖中被轉換為 $B_{i,j}$ 到 $B_{i+1,j} $ 的依賴關系）可以被如下表示：

是以，圖中這裡實線都是前向傳播時候建構的，虛線是由 fork & join 建構的。

原則上，表示虛拟依賴關系的張量可以是任意的。然而，torchgpipe選擇使用空張量，以消除由張量引起的任何不必要的計算，例如PyTorch中的梯度累積。

具體如下圖。就是使用 Fork 和 Join 的後向計算圖。圖中，不同顔色對應不同的裝置。箭頭依據後向傳播圖的方向來繪制，這些聯系是在前向傳播中被建構的。是以，$F^{'}_{i,j}$ 對于 $B_{i+1,j}$ 的虛拟依賴通過 Fork 和 Join 被建構出來，用虛線表示。

首先，我們要看看 torch.autograd.Function 的作用。

torch.autograd.Function類實際上是一個操作函數的基礎父類，這樣的操作函數必須具備兩個基本的過程，即前向的運算過程和反向的求導過程，

如果某些操作無法通過 PyTorch 已有的層或者是已有的方法實作不了，就需要實作一個新的方法對 PyTorch 進行拓展。當不使用自動求導機制，需要自定義求導規則的時候，就應該拓展torch.autograd.Function類。由于pytorch不再提供自動求導機制，就要使用者自己定義實作前向傳播和反向傳播的計算過程，這就是 "Extending torch.autograd"。

我們接下來介紹Backward Dependency 的關鍵算法：Fork and Join。

Fork 是auto grad 函數，其把一個張量 x 映射到 pair(x, $\phi$)，這裡 $\phi$ 是一個空張量。Fork 方法就是拓展了<code>torch.autograd.Function</code>。

Join 是auto grad 函數，其把 pair(x, $\phi$) 映射到一個張量 x ，這裡 $\phi$ 是一個空張量。Join 方法也是拓展了<code>torch.autograd.Function</code>。

Phony是沒有空間的張量，因為它不需要任何梯度累積，是以可以在 autograd 圖中建構任意的依賴。

在代碼中，經常可以見到 detach 的使用，這個從注釋可以看出來，是為了解決 PyTorch 的一個bug。

在 Pipeline 之中我們可以看到具體的使用方法，fence 方法（省略部分代碼）利用 depend 來建構後向傳播的依賴關系，確定 batches[i-1] 在 batches[i] 之後完成。

具體 depend 代碼如下：

我們結合示例代碼把傳入的參數指派一下，重新把方法解釋如下，這樣大家就可以更好的了解。

具體邏輯如下，通過 phony 完成了一個橋接，即在正向傳播之中，batches[i] 依賴 batches[i-1] 的執行結果。

我們把多個 batches 聯合起來看看，這樣就能看出來一個依賴鍊條。

這樣，上圖就是前向計算圖，于是在後向傳播之中，batches[i] 就必須在 batches[i-1] 之前完成了。

我們再結合論文的圖來看看。

本來示例代碼中是：

為了和論文中的圖對應，我們修改為：

depend 代碼也變化為：

對應下圖，就是在後向傳播計算圖之中 batches[i+1] 通過一個join, 一個fork，排在了 batches[i] 前面，就是下面大箭頭所示，具體細化一下：

從這個圖上，PyTorch 的 autograd 引擎不知道 $B_{i+1,j}$ 必須在 $B_{i,j}$ 之前運作，是以會打亂後向傳播的時間流。是以，虛拟依賴（前面圖的虛線箭頭）必須在前向傳播中被顯式繪制出來。

圖上的實線箭頭依據後向傳播圖的方向來繪制，這些聯系是在前向傳播中被建構的。就是說，對于 ${Batch}_i$ 來說，其反向傳播順序是固定的。就是上面一行内順序是固定的，下面一行内順序也是固定的。

但是，上下兩行之間的順序是不可知的，需要用虛線來保證，就是用 Join & Fork 來保證。

我們回頭再來看正向依賴。因為正向傳播的部分目的就是完成反向傳播依賴，而目前反向傳播隻完成了行之間的依賴，列之間的依賴沒有完成，我們現在補全。

列之間的依賴就是裝置之間的依賴，即前一個裝置的輸出是後一個裝置的輸入。

首先還是需要回顧下如何切分模型，從 split_module 可以看到，

GPipe 的 partitions 成員變量是 nn.ModuleList 類型。nn.ModuleList是一個容器，其儲存不同 module，并自動将每個 module 的 parameters 添加到網絡中。但是nn.ModuleList 并沒有定義一個網絡，而隻是将不同的子產品儲存在一起，這些子產品之間并沒有什麼先後順序，網絡的執行順序是根據 forward 函數來決定的。

随之而來問題就是：partition内部可以用Sequential來進行一系列的前向操作，但是如何配置partitions 之間的執行順序？

我們還是從論文中入手。假定我們有一個神經網絡，其由一系列子網絡構成。我們假定這些子網絡是 $f^1,...,f^n$，其參數分别是 $\theta^1,...,\theta^n$，則整個網絡是：

參數是 $\theta = (\theta^1,...,\theta^n)$，為了清楚起見，我們稱 $f^j$ 表示 f 的第 j 個分區，并假設分區的參數是互相不相交的。

在訓練網絡時，基于梯度的方法（如随機梯度下降法）需要在給定小批量訓練資料 x 和相應損失之後，計算網絡的輸出結果f(x)。以及損失相對于網絡參數 $\theta$ 的梯度g。這兩個階段分别稱為向前傳播和向後傳播。

既然 f 由其 L 層子子產品（$f^L, f^{L-1},...f^1$）順序組成，那麼前向傳播$f(x)$ 可以通過如下方式計算：讓 $x^0=x$（就是輸入x），然後順序應用每一個 partition，即 $x^j = f^j (x^{j-1})$，這裡 $ j = 1, ..., L$。就是 $f(x)$ 可以表示為 :

\[f(x) = f^L(f^{L-1}(f^{L-2}(... f^1(x))))

于是我們知道了，前向傳播的順序是由 $f(x) = f^L(f^{L-1}(f^{L-2}(... f^1(x))))$ 來确定的。

我們可以針對代碼，進一步解析，看看如何實施partitions之間的順序依賴。

解析的目标是 <code>for schedule in clock_cycles(m, n)</code> 這個 for 循環，其：

針對clock_cycles産生的每一個運作計劃：

利用 fence(schedule, skip_trackers) 建構後向傳播依賴關系。

現在我們完成了兩步：

确定性時鐘周期算法給定了前向傳播的執行順序，我們隻要按照 clock_cycles 方法提供的計劃一一運作即可。

fence 方法通過調用 join 和 fork，我們做到了在後向傳播之中，batches[i] 就必須在 batches[i-1] 之前完成了，即 $B_{i+1,j}$ 必須在 $B_{i,j}$ 之前運作。

對于我們的圖來說，第二步就是完成了下圖的列依賴。

我們的問題是：怎麼通過這個 for 循環，做到 $B_{i,{j+1}}$ 必須在 $B_{i,j}$ 之前運作？，即怎麼安排反向傳播逐次運作？就是怎麼完成行内的依賴？

這就要通過 compute 的源碼進行分析。重點說明的是：

batches[i] 這裡是會變化的，比如 batches[0] 在經過 partitions[j] 的計算之後，會變成 <code>batches[0][j]</code>。

對于 compute 方法，關鍵就是在最底部的代碼 <code>batches[i] = batch</code>。就是把第 j 個device 對第 i 個 batch 的計算結果指派到 batches[i]，指派之後，batches[i]就是 <code>batches[i][j]</code>，這樣，在下次計算時候，建構的就是 F[i, j+1], 下一次 fence 之中的 depend 操作，就是針對 <code>batches[i, j+1]</code>。

是以，在前向計算圖上，通過這個指派操作， batches[i, j+1] 就依賴 batches[i, j]，是以反向計算時候，batches[i, j + 1] 就必須在 batches[i, j] 之前完成。

關于這個指派操作，其對應的grad_fn 是 PermuteBackward，比如：

具體是：

現在，我們把下圖進行更新。

我們進行橫向拓展，得到如下，即一個batch 被分成兩個小批次: batches[i]，batches[i+1] ，它們在兩個裝置 partitions[j]，partitions[j + 1] 之上流水線，這樣行和列都有反向傳播的依賴。

手機如下：

下圖 $ m = 4, n = 3$。即，模型被分成3個子網絡，小批次被分割成 4個微批次。F 和 B 的下标是 (m, n)。

如上圖，這裡需要完成兩種依賴：

行間依賴，就是 batch 之間的依賴，就是裝置内的依賴。從圖上看是虛線，就是 $F_{1,1}$ 必須在 $F_{2,1}$之前完成，$B_{2,1}$ 必須在$B_{1,1}$ 之前完成。

列間依賴，就是 partitions（裝置）之間的依賴。從圖上看是實線，就是 $F_{1,1}$ 必須在 $F_{1,2}$之前完成，即第一個裝置必須在第二個裝置之前完成，而且第一個裝置的輸出是第二個裝置的輸入。

如上圖，我們需要完成行，列兩方面的依賴。

行間依賴是用 Join & Fork 來保證，利用空張量完成了依賴關系的設定，確定 batches[i-1] 在 batches[i] 之後完成。

列間依賴是通過 <code>batches[i] = batch</code> 完成，利用 PermuteBackward 來完成了裝置之間的依賴。

至此，我們完成了執行順序和依賴關系的設定，下一篇我們介紹如何并行處理。

一位朋友 @劍柳吟風對下面提出了疑問：

他認為：通過閱讀源碼，個人感覺應該是自定義的copy 和wait這兩個autograd.Function确定了裝置之間的依賴。

我又仔細閱讀了論文和源碼，發現他是正确的，我之前的了解有誤，特此更正如下。

我們首先看看論文内容。

論文内容如下：

2.2. Device-wise Execution Order To summarize, in pipeline parallelism (with checkpointing) each device is assigned with a set of tasks with the prescribed order. Each device will execute the given tasks one-by-one as soon as cross-device dependencies are met. However, there is a missing component in this picture data tranfer between the devices. For illustration, the full execution order that device j must follow is shown in Figure 3. Here data transfer operations are explicitly denoted as ‘receive’ and ‘send’ for emphasis.

翻譯如下：

總之，在流水線并行性（帶有檢查點）中，每個裝置都被配置設定了一組具有指定順序的任務。一旦滿足跨裝置依賴關系，每個裝置将逐個執行給定的任務。但是，在之前圖之中，裝置之間的資料傳輸中缺少一個元件。為了便于說明，裝置 j 必須遵循的完整執行順序如圖3 所示。為了更好的說明，在這裡，資料傳輸操作被明确表示為“接收”和“發送”。

具體圖例如下：

論文之中另一部分則論述了Stream 的使用。

Concurrent Copy and Computation: Streams PyTorch issues every device-bound kernels to the default stream, unless it is specified otherwise. Stream is a device- bound sequence of kernels that is executed in order. Kernels in the same stream are guaranteed to be executed in the pre- scribed order, but kernels in different streams can be inter- leaved, and even can overlap when possible. In particular, nearly all CUDA devices with compute capability 1.1 and higher support concurrent copy and execution: data transfer between devices can always overlap with kernel execution.

因為公式拷貝問題，将這段與原文下一段翻譯如下：

PyTorch将每個綁定到裝置的核釋出到預設流，除非另有規定。流是按順序執行這些綁定到裝置的核序列。同一個流中的核心保證按預先指定的順序執行，但不同流中的核可以互相交錯，甚至可能重疊。特别是，幾乎所有具有計算能力1.1及更高版本的CUDA裝置都支援并發複制和執行：裝置之間的資料傳輸總是與核心執行重疊。

torchgpipe将每個拷貝核注冊到非預設流中，同時将計算核心保留在預設流中。這允許裝置 j 可以并行處理，即 $F_{i,j}$ 可以同 "發送到裝置 $j+1$ 的 $x_{i-1}^j$" 以及/或者 "從裝置 $j-1$ 接受 $x_i^{j-1}$" 這兩個操作并行。此外，每個裝置對每個微批次使用不同的流。由于不同的微批次之間沒有真正的依賴關系，是以流的這種使用是安全的，這允許盡可能快地進行拷貝。

可見，資料傳輸是通過 Stream 來完成，即構成了實際上的裝置間依賴關系，又可以達到資料和拷貝并行的目的。

我們接下來看看具體實作，依次驗證我們的推論。

_copy_streams 定義如下：

其初始化代碼如下，chunks 是micro-batches 的數目。_ensure_copy_streams 就是針對每一個裝置的每一個macro-batch，都生成了一個專用流。

假設有3個devices，模型被分成3個子網絡，小批次被分割成 4個微批次。則具體如下：就是說 <code>_copy_streams[i][j]</code> 之中，i 表示 device 的序列，j 表示 batch 序列。（後續的文章之中，有對如何使用的詳述）

Wait 算子代碼如下，主要就是起到同步作用，等待拷貝操作的完成。

以下函數對算子進行了封裝。

fence 簡化代碼如下，其建立了圖例之中的行，列兩種依賴關系。

具體wait操作則是在 compute 之中調用，我們隻給出了部分代碼。

論文之中還有一組對比，特此翻譯摘錄：

我們還可視化了每個GPU的時間線，以幫助了解每個元件的角色，如圖所示。每幅圖檔的規劃概述如下。

(a) 通過确定性時鐘周期，所有核在前向傳播期間以正确的順序發出。時間線的左側部分對其進行了說明。然而，因為沒有明确對計算圖中依賴關系的明确編碼，autograd引擎将以無法控制的順序處理微批次，是以時間線将混亂不堪。

(b) 對于後向依賴，核心目前在向後傳播中以正确的、确定的順序釋出。

(d) Portal移除了因将跳躍張量（skipping tensor）傳輸到其間的所有裝置而導緻的不必要副本。與(c) 相比，紅色條的長度減少。

GPipe需要完成兩種依賴：

行間依賴對應了論文中的：

Pipeline parallelism’s strategy is to assign tasks with re- spect to the partition index j so that jth partition entirely lies in the jth device. In addition to this, it is enforced that Fi,j must be completed before executing Fi+1,j and Bi,j must be completed before executing Bi−1,j .

行間依賴是用 Join & Fork 來保證，利用空張量完成了依賴關系的設定，確定 batches[i-1] 在 batches[i] 之後完成。PermuteBackward 協助完成了這個依賴操作。

列間依賴是通過 Copy & Wait 兩個派生的算子來完成了裝置之間的依賴。

Markdown公式用法大全

markdown中公式編輯教程

https://docs.nvidia.com/cuda/cuda-runtime-api/stream-sync-behavior.html#stream-sync-behavior

CUDA學習：基礎知識小結

CUDA随筆之Stream的使用

NVIDIA解決方案架構師深度解析大規模參數語言模型Megatron-BERT

Accelerating Wide & Deep Recommender Inference on GPUs

HugeCTR: High-Performance Click-Through Rate Estimation Training

https://discuss.pytorch.org/t/how-to-prefetch-data-when-processing-with-gpu/548

https://github.com/NVIDIA/apex/

https://github.com/justheuristic/prefetch_generator

https://pytorch.org/tutorials/intermediate/model_parallel_turotial.html

https://pytorch.org/docs/stable/autograd.html

https://pytorch.org/docs/notes/cuda.html

https://zhuanlan.zhihu.com/p/61765561

https://pytorch.apachen.org/docs/1.7/64.html

https://zhidx.com/p/217999.html