[原始碼解析] 深度學習流水線並行之PopeDream(1)--- Profile階段

0x00 摘要

繼 GPipe 之後，我們開一個流水線並行訓練新系列，介紹微軟的PipeDream。本文介紹其總體思路，架構和Profile階段。

Gpipe流水線其存在兩個問題：硬體利用率低，記憶體佔用大。於是在另一篇流水並行的論文裡，微軟 PipeDream 針對這些問題提出了改進方法，就是1F1B （One Forward pass followed by One Backward pass）策略。這種改進策略可以解決快取 activation 的份數問題，使得 activation 的快取數量只跟 stage 數相關，從而進一步節省視訊記憶體，可以訓練更大的模型。

PipeDream 可以分為4個階段：profile，compute partition，convert model，runtime等四個階段。

Profile階段：通過小批量資料的profile推理出DNN訓練時間。
Compute Partition階段：依據profile結果確定所有層的執行時間，然後進行優化，優化器返回一個帶註釋的運算子圖，每個模型層對映到一個階段ID。
Convert model 階段：對運算子圖執行BFS遍歷，為每個階段生成一個單獨的torch.nn.Module程式碼。PipeDream對每個階段中的運算子進行排序，以確保它們保持與原始PyTorch模型圖的輸入輸出依賴關係的一致性。
Runtime 階段：PipeDream執行時根據其1F1B-RR排程策略將每個階段（包括複製階段的副本）分配給單個工作程式。

本文首先看看 PipeDream 總體思路，架構和Profile階段。

0x01 概述

1.1 前文回顧

前文提到，目前分散式模型訓練有幾個必要並行技術：

流水線並行，尤其是如何自動設定流水；
梯度累加（Gradient Accumulation）；
後向重計算；
1F1B 策略；

在前面幾篇文章之中，我們介紹了Gpipe 如何前三種技術。本文開始，我們通過微軟分散式DNNs訓練系統PipeDream來看看其如何實現流水線並行和1F1B 策略。

1.2 目前問題

DNN 訓練的特點是雙向訓練，訓練在前向和後向通道中迭代進行計算，兩種傳播以相反順序穿過相同層，在每輪迭代中，訓練過程迴圈處理輸入資料的一個 minibatch，並且更新模型引數。

1.2.1 資料並行

最常見的 DNN 並行化訓練方法是資料並行化，這種方法把輸入資料分散到各個 workers中（每個worker擁有全部模型）執行。不幸的是，儘管在加速資料並行的效能優化方面取得了一些進展，但若要在雲基礎設施上訓練就會產生很大的通訊開銷。而且，隨著GPU 計算速度的飛快提升，所有模型訓練的耗時瓶頸會更進一步轉向通訊環節。

下圖為資料並行，但是其對硬體利用率依然太低。因為資料並行化中的單個 worker 在交換梯度資料時不得不進行通訊等待。

1.2.2 模型並行

模型並行化是另一種並行化訓練形式，這種方法是把運算元分散到各個 worker 上進行計算，這通常用於訓練大型 DNN 模型。本文中，模型並行指的就是把模型的不同layer放在不同的機器上，不涉及將同一個layer切分到不同機器上的場景。

下圖是模型並行化，顯示了一個計算時間線，該示例有四臺機器和一個管道。

在正向階段，每個階段對本階段中的層的minibatch執行正向傳遞，並將結果傳送到下一階段。輸出級在完成前向傳遞後，計算minibatch的損失。
在後向階段，每個階段形成後向通道，逐一將損失傳播到前一階段。

worker 之間只能同時處理一個 minibatch，系統中只有一個minibatch是活動的，這極大限制了硬體利用率。

而且，模型並行化需要程式設計師決定怎樣按照給定的硬體資源來分割特定的模型，這其實無形之中加重了程式設計師的負擔。

1.2.3 Gpipe

除了通用問題之外，GPipe 流水並行策略還有一個記憶體問題：需要快取多份activation。

假如一個batch被分為 n 個micro-batch，則需要快取 n 份activation。這個 n 是梯度累加的次數，為了儘可能流水，通常這個累加次數都比較大，一般大於兩倍 stage 數目。那麼即使只快取少數 Tensor，這種策略依然需要較多視訊記憶體。

0x02 論文

PipeDream 就針對這些問題提出了改進方法1F1B。PipeDream是第一個以自動化和通用的方式將流水線並行，模型並行和資料並行結合起來的系統。PipeDream首先使用模型並行對DNN進行劃分，並將每層的子集分配給每個worker。但是與傳統的模型並行不同，PipeDream對小批量資料進行流水線處理，實現了潛在的管道並行設計。在任何時刻，不同的worker處理不同的輸入，從而保證了流水線的滿負荷以及並行BSP。

微軟在論文 PipeDream: Fast and Efficient Pipeline Parallel DNN Training 之中對於PipeDream進行了詳細闡述，所以我們就基於此論文進行分析。

2.1 方案概述

2.1.1 並行方式

PipeDream 模型的基本單位是層，PipeDream將DNN的這些層劃分為多個階段。每個階段（stage）由模型中的一組連續層組成。

PipeDream的主要並行方式就是把模型的不同層放到不同的stage之上，不同的的stage部署在不同的機器上，順序地進行前向和反向計算，形成了一個pipeline。

每個階段對該階段中的所有層執行向前和向後傳遞。PipeDream將包含輸入層的階段稱為輸入階段，將包含輸出層的階段稱為輸出階段。但是每個stage可能有不同的replication，這就是資料並行。

對於使用資料並行的stage，採用 round-robin的方式將任務分配在各個裝置上，需要保證一個batch的資料在前向和後向發生在同一臺機器上，

2.1.2 1F1B

由於前向計算的 activation 需要等到對應的後向計算完成後才能釋放（無論有沒有使用 Checkpointing 技術），因此在流水並行下，如果想盡可能節省快取 activation 的份數，就要儘量縮短每份 activation 儲存的時間，也就是讓每份 activation 都儘可能早的釋放，所以要讓每個 micro-batch 的資料儘可能早的完成後向計算，因此需要把後向計算的優先順序提高，讓 micro-batch 標號小的後向比 micro-batch 標號大的前向先做。因此，如果我們讓最後一個 stage 在做完一次 micro-batch 的前向後，立馬就做本 micro-batch 的後向，那麼我們就能讓其他的 stage 儘可能早的開始後向計算，這就是 1F1B 策略。

1F1B（one-forward-one-backward）的排程模式會在每臺worker機器上交替進行小批次資料的前向後向計算，同時確保這些小批量在"後向傳播"時可以路由到"前向傳播"的相同worker。

這種方案可以使得每個GPU上都會有一個batch的資料正在被處理，使所有worker保持忙碌，而不會出現管道暫停，整個pipeline是比較均衡的。同時能確保以固定週期執行每個stage上的引數更新，也有助於防止同時處理過多小批量並確保模型收斂。

0x03 流水線

PipeDream的pipeline parallelism（PP）是一種新的並行化策略，它將批內並行與批間並行結合起來。

3.1 流水線改進

我們首先看看流水線對於模型並行的改進。

在上節的示例中，如果只有一個活動的minibatch，系統在任何給定的時間點最多有一個GPU處於活動狀態。

但是我們希望所有GPU都處於活動狀態。因此PipeDream 多個小批量逐一注入到流水線中，從而通過流水線來增強模型並行訓練。在完成小批量的前向傳遞時，每個階段都會非同步地將輸出啟用傳送到下一階段，同時開始處理另一個小批量。類似地，在完成一個小批量的向後傳遞後，每個階段都會將梯度非同步傳送到前一階段，同時開始計算另一個小批量。

與普通層間並行訓練相比，流水線有兩個主要優點：

流水線通訊量較少。流水線並行比資料並行的通訊量要少得多。與資料並行方法（使用集體通訊或引數伺服器）中的做法（把所有引數聚合梯度並且將結果傳送給所有worker）不同，流水線並行中的每個worker只需要在兩個 stage 邊界之間將梯度和輸出啟用的一個子集傳送給另一個worker，這可能導致某些模型的通訊量大幅減少。
流水線重疊了計算和通訊。跨階段前向輸出啟用和後向梯度的非同步通訊可以使得這些通訊與後續小批量計算在時間上重疊，因為它們在不同的輸入上執行，計算和通訊完全獨立，沒有依賴邊，所以可以更容易的並行化。在穩定理想狀態下，所有的 workers 時刻都在運轉，不像模型並行化訓練中會有停下來等待的時候。

下圖是實施了 1F1B 的流水線。Machine 1先計算藍色 1，然後把藍色 1 傳送給 Machine 2 繼續計算。Machine 1 接著計算藍色 2。Machine 1 和 Machine 2 之間只傳送模型的一個子集。計算和通訊可以並行。

3.2 挑戰

PipeDream的目標是以最小化總體訓練時間的方式將流水線並行，模型並行性和資料並行性結合起來。然而，要使這種方法對大型DNN模型有效，獲得流水線並行化訓練的潛在收益，PipeDream 必須克服幾個主要挑戰：

如何高效劃分流水線。與處理器中的流水線一樣，需要將DNN高效正確地劃分為若干“階段”（層序列），每個階段部署在不同的worker上執行。
- 模型特質和硬體拓撲會降低效率，劃分應該具體取決於模型體系結構和硬體部署。不好的劃分（階段性的工作量大範圍傾斜）可能會導致worker長時間閒置。所以需要依據一定原則（通訊和資源利用率）來劃分，比如：彼此有通訊的層應該分配到相鄰的處理器；如果多個層操作同一資料結構，它們應該被分配到同一個處理器上，彼此獨立的層可以對映到不同處理器上。所以分配演算法也必須考慮模型特質和硬體拓撲。
- 機器間的過度通訊會降低硬體效率。
- 在確保訓練任務向前推進的同時，如何排程計算以最大化吞吐量。
如何防止流水線瓶頸。
- 由木桶原理我們可以知道，在穩定狀態下，一個流水線管道的吞吐量由這個流水線上最慢環節的吞吐量決定。如果各個環節的處理能力彼此差別很大，會導致管道中出現空閒時間（一半稱之為bubble），這樣最快環節必須停下來等待其他環境，會造成飢餓現象，從而導致資源利用不足。所以需要確保流水線中所有階段都大致花費相同的計算時間，否則最慢的階段將會成為整個流水線的瓶頸。
如何在不同的輸入資料之間排程工作以均衡流水線。
- 與傳統的單向流水線管道不同，DNN訓練是雙向的：前向傳播和後向傳播，兩種傳播以相反順序穿過相同層。如何協調流水線工作是一個問題。
面對流水線帶來的非同步性，如何確保訓練有效。
- 流水線帶來的一個問題就是weight版本眾多。在後向傳播時候如果使用比前向傳播時更高版本的weight來計算，則會造成訓練模型質量降低。
- PipeDream管理後向通道里的權重版本，通過為每個小批量的weight維護版本號來解決這個問題，這樣在後向通道里使用的權重版本就和前向通道里使用的相同，從而在數值上能夠正確計算梯度（我們後續文章會講解）。

3.4 流水線劃分演算法

PipeDream基於一個短期執行分析結果來自動劃分DNN的層，依據分析結果使用演算法來對不同階段之間的計算負載進行平衡，同時最小化通訊。PipeDream的自動劃分演算法總體目標是輸出一個平衡的管道，確保每個階段大致執行相同的總工作量。同時還必須確保各階段之間通訊的資料量儘可能小，以避免通訊中斷。演算法如下：

將DNN層劃分為多個階段，以便每個階段以大致相同的速率完成，即花費大致相同的計算時間。
嘗試以拓撲感知的方式儘量減少worker之間的通訊（例如，如果可能，向更高頻寬的鏈路傳送較大的輸出）。
因為DNN並不總可以在可用的workers做平均分配，為了進一步改進負載平衡，PipeDream允許複製一個stage，即在這個stage上使用多個worker進行資料並行。這樣多個worker可以分配到流水線同一階段，並行處理一個batch的不同的mini-batch，提高處理效率。因為資料並行採用了RR，所以這套策略也被稱為 1F1B-RR(one-forward-noe-backward-round-robin)。

這個劃分問題等價於最小化流水線的最慢階段所花費的時間，並且具有最優子問題屬性：在給定worker工作量前提下，吞吐量最大化的流水線由一系列子流水線構成，其中每一個子流水線針對較小worker工作量來最大化自己的輸出。因此，PipeDream使用動態規劃來尋找最優解。

具體如下圖：

3.5 Profile

DNN訓練有一個特點：不同輸入的計算時間幾乎沒有變化。於是 PipeDream充分利用了這一事實，給定一個具有N層和M臺可用機器的DNN，PipeDream首先在一臺機器上分析模型，記錄向前和向後過程所花費的計算時間，層輸出的大小以及每個層的相關引數的大小，最後輸出為一個結果檔案。

分割槽演算法不但使用profile結果檔案作為輸入，而且還考慮了其他限制，如硬體拓撲和頻寬、工人數量和計算裝置的記憶體容量，最終將層分為多個階段，同時還確定每個階段的複製因子，以最小化模型的總訓練時間。

所以總體演算法大致如下：

因為PipeDream借鑑了很多GPipe的思路，所以可以看到其比Gpipe的進步之處。

比如 Gpipe是通過在程式碼中硬性預估ops來進行流水線負載均衡，PipeDream則是先做profile，根據實際情況再做推理。

0x04 Profile階段

Profile是PipeDream工作的第一個階段，是分割槽演算法的基礎。PipeDream根據profiling的結果，使用動態規劃對模型進行劃分，將模型劃分為不同的stage，以及每個stage的replication數。

這是PipeDream針對GPipe的一個改進，兩者都是對每層的執行時間進行預估，然後對模型進行劃分。

GPipe是利用經驗值或者數學的方法來對執行時間進行預估。
PipeDream根據profiling的結果對執行時間進行預估。

因為有實際資料進行支撐，所以PipeDream更加準確和先進。

4.1 思路

評測機制利用了這樣一個事實：DNN訓練在計算和通訊時間上幾乎沒有變化。所以我們可以通過小批量資料的profile推理出DNN訓練時間。為了確定所有層的執行時間，PipeDream在其中一臺機器上使用1000個小批量對DNN模型的短期（幾分鐘）執行進行 profile。

4.1.1 如何計算

執行時間

對於每一層的執行時間，我們可以通過執行時間 = 計算時間 + 通訊時間來得到。

計算時間就是每層layer前向和後向的計算時間，這個可以從profile得出。
通訊時間就需要根據模型大小進行估算，PipeDream 估計通訊所需的時間為"需要傳輸的資料量"除以"通訊鏈路上的頻寬"。

通訊時間

在流水線上，大多數通訊都有三個步驟：

1）在傳送端機器上，從GPU傳輸到CPU移動資料。

2）通過網路從傳送者到接收者傳送資料。

3）在接收端，從CPU到GPU移動資料。

而通過網路從傳送者到接收者傳送資料是最耗時的，所以PipeDream主要考慮這個因素。如果再對這個因素細分，則有：

對於從層 i 到層 i + 1 傳輸啟用值的時間，PipeDream 基於 "啟用值"來估計。
假如配置成了資料並行（對於層 i 使用 m 個 worker 做資料並行）的情況，做權重同步的時間使用"權重"來估計：
- 如果使用分散式引數伺服器，則權重數量被預估為 4 x ( m - 1 ) x | w i | / m。
- 如果使用 all_reduce，則每個worker給其他workers傳送 ( m - 1 ) x | w i | / m 個 bytes，也接受到同樣數量位元組。

4.1.2 Profile內容

綜上所述，PipeDream在profile之中，為每個層 i 記錄三個數量：

Ti，層 i 的在GPU上向前和向後計算時間之和，即每層layer前向和後向的計算時間；
ai，層 i 的輸出啟用的大小（以及向後過程中輸入梯度的大小）以位元組為單位，即每層layer的輸出的大小；
wi，層 i 的權重引數的大小（以位元組為單位），即每層layer引數的大小；

4.2 程式碼

不同模型或者說不同領域有不同的profile檔案。

我們以 profiler/translation/train.py 為入口進行分析。

4.2.1 訓練指令碼

以下我們省略了無關程式碼。

4.2.1.1 訓練過程

class Seq2SeqTrainer:

    def feed_data(self, data_loader, training=True):
        """
        Runs training or validation on batches from data_loader.

        :param data_loader: data loader
        :param training: if True runs training else runs validation
        """
        # 白名單
        module_whitelist = ["EmuBidirLSTM", "RecurrentAttention", "Classifier"]
        
        # 樣本集
        for i, (src, tgt) in enumerate(data_loader):
            break
        (src, src_length) = src
        (tgt, tgt_length) = tgt
        src_length = torch.LongTensor(src_length).cuda()
        src = src.cuda()
        tgt = tgt.cuda()
        model_input = (src, src_length, tgt[:-1])
        # 使用torchsummary計算網路的計算引數等資訊
        summary = torchsummary.summary(model=self.model, module_whitelist=module_whitelist,
                                       model_input=model_input, verbose=True)

         for i, (src, tgt) in enumerate(data_loader):

            if training and i in eval_iters:
                test_bleu, _ = self.translator.run(calc_bleu=True,
                                                   epoch=self.epoch,
                                                   iteration=i)
                # 訓練模型
                self.model.train()
                self.preallocate(data_loader, training=True)
                        
        # 從模型建立圖      
        if training:
            create_graph(self.model, module_whitelist, (src, tgt), summary,
                         os.path.join("profiles", self.arch))

4.2.1.2 計算引數

上節在訓練指令碼製作的時候，torchsummary 的作用是計算網路的計算引數等資訊，對於 torchsummary 我們舉例如下：

import torch
import torch.nn as nn
from torchsummary import summary

class SimpleConv(nn.Module):
    def __init__(self):
        super(SimpleConv, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 1, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
        )

    def forward(self, x, y):
        x1 = self.features(x)
        x2 = self.features(y)
        return x1, x2
    
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = SimpleConv().to(device)

summary(model, [(1, 16, 16), (1, 28, 28)])

其列印如下：

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1            [-1, 1, 16, 16]              10
              ReLU-2            [-1, 1, 16, 16]               0
            Conv2d-3            [-1, 1, 28, 28]              10
              ReLU-4            [-1, 1, 28, 28]               0
================================================================
Total params: 20
Trainable params: 20
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.77
Forward/backward pass size (MB): 0.02
Params size (MB): 0.00
Estimated Total Size (MB): 0.78
----------------------------------------------------------------

4.2.1.3 建立圖

create_graph 的作用就是使用torchgraph.GraphCreator建立一個圖，這個圖就可以理解為模型內部的DAG圖，每個節點記錄如下資訊。

node10 -- Dropout(p=0.2) -- forward_compute_time=0.064, backward_compute_time=0.128, activation_size=6291456.0, parameter_size=0.000

具體程式碼如下：

def create_graph(model, module_whitelist, model_input, summary, directory):
    """Given a model, creates and visualizes the computation DAG
       of the model in the passed-in directory."""
    # 建立圖
    graph_creator = torchgraph.GraphCreator(model, summary, module_whitelist
    # 構建hook                                        
    graph_creator.hook_modules(model, root=True) 
    (src, tgt) = model_input
    (src, src_length) = src
    (tgt, tgt_length) = tgt
    src_length = torch.LongTensor(src_length).cuda()
    src = src.cuda()
    tgt = tgt.cuda()
    # 執行以得到profile                                        
    model(src, src_length, tgt[:-1])
    graph_creator.unhook_modules()
    # 輸出profile結果                                        
    graph_creator.persist_graph(directory)

4.2.2 建立圖

建立圖基本是在 GraphCreator 內完成。

class GraphCreator(object):
    def __init__(self, model, summary, module_whitelist):
        if isinstance(model, torch.nn.Module) is False:
            raise Exception("Not a valid model, please provide a 'nn.Module' instance.")

        self.model = model
        self.module_whitelist = module_whitelist
        self.summary = copy.deepcopy(summary)
        self.forward_original_methods = {}
        self.graph = graph.Graph()
        self.inputs = {}

4.2.2.1 設定wrapper

hook_modules 的作用是給模型的forward函式設定一個wrapper，並且遍歷為子模組設定，這樣在模型執行時候可以跟蹤模型之間的聯絡。

def hook_modules(self, module, root=False):
    this_creator = self
    sub_modules = module.__dict__['_modules']

    # Wrapper function to "forward()", keeping track of dependencies.
    def forward_wrapper(self, *wrapped_inputs):
        input = []
        wrapped_inputs_list = list(wrapped_inputs)
        for i in range(len(wrapped_inputs_list)): # 遍歷輸入
            if isinstance(wrapped_inputs_list[i], TensorWrapper):
                # 如果已經被包裝，則插入input
                input.append(wrapped_inputs_list[i].tensor)
            else:
                key = wrapped_inputs_list[i]
                if key in this_creator.inputs: # 如果是原始輸入，則不進行包裝
                    wrapped_inputs_list[i] = this_creator.inputs[key]
                else:
                    j = len(this_creator.inputs)
                    # 如果沒有被wrap, 則構建一個TensorWrapper進行包裝
                    wrapped_inputs_list[i] = TensorWrapper(wrapped_inputs_list[i],
                                                           "Input%d" % j, this_creator)
                    this_creator.inputs[key] = wrapped_inputs_list[i]
                input.append(wrapped_inputs_list[i].tensor) # 則插入input
        result = this_creator.forward_original_methods[self](*input)
        # 對結果進行包裝
        wrapped_result = TensorWrapper(result, str(self), this_creator)
        
        # 把邊新增進入圖
        for wrapped_input in wrapped_inputs_list:
            this_creator.graph.add_edge(wrapped_input.node(), wrapped_result.node())

        return wrapped_result

    # Wrapper function to "forward()", keeping track of dependencies.
    def forward_wrapper_root(self, *wrapped_inputs):
        input = []
        wrapped_inputs_list = list(wrapped_inputs)
        for i in range(len(wrapped_inputs_list)):
            if isinstance(wrapped_inputs_list[i], TensorWrapper):
                input.append(wrapped_inputs_list[i].tensor)
            else:
                key = wrapped_inputs_list[i]
                if key in this_creator.inputs:
                    wrapped_inputs_list[i] = this_creator.inputs[key]
                else:
                    j = len(this_creator.inputs)
                    wrapped_inputs_list[i] = TensorWrapper(wrapped_inputs_list[i],
                                                           "Input%d" % j, this_creator)
                    this_creator.inputs[key] = wrapped_inputs_list[i]
                input.append(wrapped_inputs_list[i].tensor)
        result = this_creator.forward_original_methods[self](*input)

        return result

    # 遍歷子模組，遞迴設定wrapper  
    for name, sub_module in sub_modules.items():
        # nn.Module is the only thing we care about.
        if sub_module is None or isinstance(sub_module, torch.nn.Module) is False:
            break

        sub_module_name = sub_module.__class__.__name__
        sub_sub_modules = sub_module.__dict__['_modules']
        if len(sub_sub_modules) == 0 or sub_module_name in self.module_whitelist:
            sub_module.reset_hooks()
            #
            # Hook nn.Module with no descendants.
            #

            # Replace "forward" with "wrapped_forward".
            # 使用wrapped_forward替換forward
            if sub_module not in this_creator.forward_original_methods:
                this_creator.forward_original_methods.update({sub_module:
                                                               sub_module.forward})
                sub_module.forward = forward_wrapper.__get__(sub_module, sub_module.__class__)

        if len(sub_sub_modules) >forward_compute_time 0 and sub_module_name not in self.module_whitelist:
            #
            # Recursively visit this module's descendants.
            # 遞迴設定wrapper
            self.hook_modules(sub_module)
    if root: # 對於root進行處理
        this_creator.forward_original_methods.update({module: module.forward})
        module.forward = forward_wrapper_root.__get__(module, module.__class__)

4.2.2.2 TensorWrapper

TensorWrapper 就實現了wrapper功能，graph_creator.summary 就是之前torchsummary.summary得到的網路等資訊。可以看到此類會遍歷 summary，計算 forward_compute_time 等資訊，最終構建了一個 node。

需要注意的是：activation_sizes 是根據 output_shape 來計算的。

class TensorWrapper(object):
    def __init__(self, tensor, node_desc, graph_creator, activation_size=None):
        self.tensor = tensor
        global object_id
        self.object_id = object_id
        object_id += 1
        self.node_desc = node_desc

        i = 0
        for i in range(len(graph_creator.summary)):
            if str(graph_creator.summary[i]['layer_name']) == node_desc:
                break

        if i < len(graph_creator.summary) and node_desc == str(graph_creator.summary[i]['layer_name']):
            summary_elem = graph_creator.summary.pop(i)
            forward_compute_time = summary_elem['forward_time']
            backward_compute_time = summary_elem['backward_time']
            if isinstance(summary_elem['output_shape'][0], list):
                activation_sizes = [4.0 * functools.reduce(lambda x, y: x * y, elem)
                                    for elem in summary_elem['output_shape']]
            else:
                activation_sizes = 4.0 * functools.reduce(lambda x, y: x * y, summary_elem['output_shape'])
            parameter_size = 4.0 * float(summary_elem['nb_params'])
            self._node = graph.Node("node%d" % object_id, node_desc=node_desc,
                                    forward_compute_time=forward_compute_time,
                                    backward_compute_time=backward_compute_time,
                                    activation_size=activation_sizes,
                                    parameter_size=parameter_size)
        elif activation_size is not None:
            self._node = graph.Node("node%d" % object_id, node_desc=node_desc,
                                    activation_size=activation_size)
        else:
            self._node = graph.Node("node%d" % object_id, node_desc=node_desc)
        self.graph_creator = graph_creator

對於某些內建方法，則也會相應處理，比如如下。

def __iadd__(self, other):
    self_activation_size = self.node().activation_size
    other_activation_size = other.node().activation_size
    assert(self_activation_size == other_activation_size)
    wrapped_result = TensorWrapper(self.tensor, "Add(inplace)", self.graph_creator,
                                   activation_size=self_activation_size)
    self.tensor += other.tensor
    self.graph_creator.graph.add_edge(self._node, wrapped_result.node())
    self.graph_creator.graph.add_edge(other.node(), wrapped_result.node())
    return wrapped_result

最終對應：

node58 -- Add(inplace) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=102760448.000, parameter_size=0.000

4.2.3 持久化

persist_graph 就是把profile結果輸出到檔案。

def persist_graph(self, directory):
    self.graph.to_dot(os.path.join(directory, "graph.dot"))
    with open(os.path.join(directory, "graph.txt"), 'w') as f:
        f.write(str(self.graph))
    self.graph.render_bar_graphs_and_cdfs(directory)

具體呼叫了 graph.py 的函式完成，這裡摘錄 to_dot函式如下：

def to_dot(self, arch):
    dot = graphviz.Digraph()
    for node in self.nodes.values():
        node_desc = "%s\n[forward_compute_time=%.3f,backward_compute_time=%.3f,activation_size=%s,parameter_size=%.1f]" % (
            node.node_desc, node.forward_compute_time, node.backward_compute_time,
            node.activation_size, node.parameter_size)
        if node.stage_id is not None:
            color = self._colors[node.stage_id % len(self._colors)]
            dot.node(node.node_id, node_desc,
               color=color, style='filled')
        else:
            dot.node(node.node_id, node_desc)
    for node in self.nodes.values():
        if node.node_id not in self.edges:
            continue
        for out_node in self.edges[node.node_id]:
            dot.edge(node.node_id, out_node.node_id)
    dot.render(arch)

4.3 結果

我們使用原始碼中的結果為例 pipedream-pipedream/profiler/translation/profiles/gnmt/graph.txt，給大家展示下具體結果。

node1 -- Input0 -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=0.0, parameter_size=0.000
node4 -- Embedding(32320, 1024, padding_idx=0) -- forward_compute_time=0.073, backward_compute_time=6.949, activation_size=6291456.0, parameter_size=132382720.000
node5 -- EmuBidirLSTM(  (bidir): LSTM(1024, 1024, bidirectional=True)  (layer1): LSTM(1024, 1024)  (layer2): LSTM(1024, 1024)) -- forward_compute_time=5.247, backward_compute_time=0.016, activation_size=12582912.0, parameter_size=67174400.000
node2 -- Input1 -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=0.0, parameter_size=0.000
node6 -- Dropout(p=0.2) -- forward_compute_time=0.077, backward_compute_time=0.196, activation_size=12582912.0, parameter_size=0.000
node7 -- LSTM(2048, 1024) -- forward_compute_time=3.190, backward_compute_time=5.348, activation_size=[6291456.0; 131072.0; 131072.0], parameter_size=50364416.000
node8 -- __getitem__(0) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6291456.0, parameter_size=0.000
node9 -- __getitem__(1) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=131072.0, parameter_size=0.000
node10 -- Dropout(p=0.2) -- forward_compute_time=0.064, backward_compute_time=0.128, activation_size=6291456.0, parameter_size=0.000
node11 -- LSTM(1024, 1024) -- forward_compute_time=2.491, backward_compute_time=4.203, activation_size=[6291456.0; 131072.0; 131072.0], parameter_size=33587200.000
node12 -- __getitem__(0) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6291456.0, parameter_size=0.000
node13 -- __getitem__(1) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=131072.0, parameter_size=0.000
node14 -- Add -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6291456.0, parameter_size=0.000
node15 -- Dropout(p=0.2) -- forward_compute_time=0.059, backward_compute_time=0.121, activation_size=6291456.0, parameter_size=0.000
node16 -- LSTM(1024, 1024) -- forward_compute_time=2.492, backward_compute_time=4.201, activation_size=[6291456.0; 131072.0; 131072.0], parameter_size=33587200.000
node17 -- __getitem__(0) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6291456.0, parameter_size=0.000
node18 -- __getitem__(1) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=131072.0, parameter_size=0.000
node19 -- Add -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6291456.0, parameter_size=0.000
node3 -- Input2 -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=0.0, parameter_size=0.000
node21 -- Embedding(32320, 1024, padding_idx=0) -- forward_compute_time=0.066, backward_compute_time=0.328, activation_size=6291456.0, parameter_size=132382720.000
node20 -- hidden -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=0.0, parameter_size=0.000
node22 -- __getitem__(0) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=0.0, parameter_size=0.000
node23 -- RecurrentAttention(  (rnn): LSTM(1024, 1024)  (attn): BahdanauAttention(    (linear_q): Linear(in_features=1024, out_features=1024, bias=False)    (linear_k): Linear(in_features=1024, out_features=1024, bias=False)    (dropout): Dropout(p=0)  )  (dropout): Dropout(p=0)) -- forward_compute_time=4.546, backward_compute_time=6.141, activation_size=[6160384.0; 131072.0; 131072.0; 6160384.0; 288768.0], parameter_size=41979904.000
node24 -- __getitem__(0) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6160384.0, parameter_size=0.000
node25 -- __getitem__(1) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=131072.0, parameter_size=0.000
node26 -- __getitem__(2) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=131072.0, parameter_size=0.000
node27 -- __getitem__(3) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6160384.0, parameter_size=0.000
node28 -- Dropout(p=0.2) -- forward_compute_time=0.058, backward_compute_time=0.176, activation_size=6160384.0, parameter_size=0.000
node29 -- Concat(2) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6160384.0, parameter_size=0.000
node30 -- __getitem__(1) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=0.0, parameter_size=0.000
node31 -- LSTM(2048, 1024) -- forward_compute_time=3.151, backward_compute_time=5.288, activation_size=[6160384.0; 131072.0; 131072.0], parameter_size=50364416.000
node32 -- __getitem__(0) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6160384.0, parameter_size=0.000
node33 -- __getitem__(1) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=131072.0, parameter_size=0.000
node34 -- Dropout(p=0.2) -- forward_compute_time=0.061, backward_compute_time=0.174, activation_size=6160384.0, parameter_size=0.000
node35 -- Concat(2) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6160384.0, parameter_size=0.000
node36 -- __getitem__(2) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=0.0, parameter_size=0.000
node37 -- LSTM(2048, 1024) -- forward_compute_time=3.145, backward_compute_time=5.306, activation_size=[6160384.0; 131072.0; 131072.0], parameter_size=50364416.000
node38 -- __getitem__(0) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6160384.0, parameter_size=0.000
node39 -- __getitem__(1) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=131072.0, parameter_size=0.000
node40 -- Add -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6160384.0, parameter_size=0.000
node41 -- Dropout(p=0.2) -- forward_compute_time=0.055, backward_compute_time=0.198, activation_size=6160384.0, parameter_size=0.000
node42 -- Concat(2) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6160384.0, parameter_size=0.000
node43 -- __getitem__(3) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=0.0, parameter_size=0.000
node44 -- LSTM(2048, 1024) -- forward_compute_time=3.149, backward_compute_time=15.883, activation_size=[6160384.0; 131072.0; 131072.0], parameter_size=50364416.000
node45 -- __getitem__(0) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6160384.0, parameter_size=0.000
node46 -- __getitem__(1) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=131072.0, parameter_size=0.000
node47 -- Add -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6160384.0, parameter_size=0.000
node48 -- Classifier(  (classifier): Linear(in_features=1024, out_features=32320, bias=True)) -- forward_compute_time=5.609, backward_compute_time=1.227, activation_size=194437120.0, parameter_size=132512000.000
   node1 -- node4
   node4 -- node5
   node2 -- node5
   node5 -- node6
   node6 -- node7
   node7 -- node8
   node7 -- node9
   node8 -- node10
   node10 -- node11
   node11 -- node12
   node11 -- node13
   node12 -- node14
   node8 -- node14
   node14 -- node15
   node15 -- node16
   node16 -- node17
   node16 -- node18
   node17 -- node19
   node14 -- node19
   node3 -- node21
   node20 -- node22
   node21 -- node23
   node22 -- node23
   node19 -- node23
   node2 -- node23
   node23 -- node24
   node23 -- node25
   node23 -- node26
   node23 -- node27
   node24 -- node28
   node28 -- node29
   node26 -- node29
   node20 -- node30
   node29 -- node31
   node30 -- node31
   node31 -- node32
   node31 -- node33
   node32 -- node34
   node34 -- node35
   node26 -- node35
   node20 -- node36
   node35 -- node37
   node36 -- node37
   node37 -- node38
   node37 -- node39
   node38 -- node40
   node32 -- node40
   node40 -- node41
   node41 -- node42
   node26 -- node42
   node20 -- node43
   node42 -- node44
   node43 -- node44
   node44 -- node45
   node44 -- node46
   node45 -- node47
   node40 -- node47
   node47 -- node48

至此，我們知道了Profile階段的內容，就是：執行訓練指令碼，依據執行結果來計算引數，建立一個模型內部的DAG圖，然後把引數和DAG圖持久化到檔案之中，後續階段會使用這個檔案的內容。

下一篇我們分析如何計算自動分割槽。