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

0x00 摘要

繼 GPipe 之后，我們開(kāi)一個(gè)流水線并行訓(xùn)練新系列，介紹微軟的PipeDream。本文介紹其總體思路，架構(gòu)和Profile階段。

Gpipe流水線其存在兩個(gè)問(wèn)題：硬件利用率低，內(nèi)存占用大。于是在另一篇流水并行的論文里，微軟 PipeDream 針對(duì)這些問(wèn)題提出了改進(jìn)方法，就是1F1B （One Forward pass followed by One Backward pass）策略。這種改進(jìn)策略可以解決緩存 activation 的份數(shù)問(wèn)題，使得 activation 的緩存數(shù)量只跟 stage 數(shù)相關(guān)，從而進(jìn)一步節(jié)省顯存，可以訓(xùn)練更大的模型。

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

Profile階段：通過(guò)小批量數(shù)據(jù)的profile推理出DNN訓(xùn)練時(shí)間。
Compute Partition階段：依據(jù)profile結(jié)果確定所有層的運(yùn)行時(shí)間，然后進(jìn)行優(yōu)化，優(yōu)化器返回一個(gè)帶注釋的運(yùn)算符圖，每個(gè)模型層映射到一個(gè)階段ID。
Convert model 階段：對(duì)運(yùn)算符圖執(zhí)行BFS遍歷，為每個(gè)階段生成一個(gè)單獨(dú)的torch.nn.Module代碼。PipeDream對(duì)每個(gè)階段中的運(yùn)算符進(jìn)行排序，以確保它們保持與原始PyTorch模型圖的輸入輸出依賴關(guān)系的一致性。
Runtime 階段：PipeDream運(yùn)行時(shí)根據(jù)其1F1B-RR調(diào)度策略將每個(gè)階段（包括復(fù)制階段的副本）分配給單個(gè)工作進(jìn)程。

本文首先看看 PipeDream 總體思路，架構(gòu)和Profile階段。

0x01 概述

1.1 前文回顧

前文提到，目前分布式模型訓(xùn)練有幾個(gè)必要并行技術(shù)：

流水線并行，尤其是如何自動(dòng)設(shè)定流水；
梯度累加（Gradient Accumulation）；
后向重計(jì)算；
1F1B 策略；

在前面幾篇文章之中，我們介紹了Gpipe 如何前三種技術(shù)。本文開(kāi)始，我們通過(guò)微軟分布式DNNs訓(xùn)練系統(tǒng)PipeDream來(lái)看看其如何實(shí)現(xiàn)流水線并行和1F1B 策略。

1.2 目前問(wèn)題

DNN 訓(xùn)練的特點(diǎn)是雙向訓(xùn)練，訓(xùn)練在前向和后向通道中迭代進(jìn)行計(jì)算，兩種傳播以相反順序穿過(guò)相同層，在每輪迭代中，訓(xùn)練過(guò)程循環(huán)處理輸入數(shù)據(jù)的一個(gè) minibatch，并且更新模型參數(shù)。

1.2.1 數(shù)據(jù)并行

最常見(jiàn)的 DNN 并行化訓(xùn)練方法是數(shù)據(jù)并行化，這種方法把輸入數(shù)據(jù)分散到各個(gè) workers中（每個(gè)worker擁有全部模型）運(yùn)行。不幸的是，盡管在加速數(shù)據(jù)并行的性能優(yōu)化方面取得了一些進(jìn)展，但若要在云基礎(chǔ)設(shè)施上訓(xùn)練就會(huì)產(chǎn)生很大的通信開(kāi)銷。而且，隨著GPU 計(jì)算速度的飛快提升，所有模型訓(xùn)練的耗時(shí)瓶頸會(huì)更進(jìn)一步轉(zhuǎn)向通信環(huán)節(jié)。

下圖為數(shù)據(jù)并行，但是其對(duì)硬件利用率依然太低。因?yàn)閿?shù)據(jù)并行化中的單個(gè) worker 在交換梯度數(shù)據(jù)時(shí)不得不進(jìn)行通信等待。

1.2.2 模型并行

模型并行化是另一種并行化訓(xùn)練形式，這種方法是把算子分散到各個(gè) worker 上進(jìn)行計(jì)算，這通常用于訓(xùn)練大型 DNN 模型。本文中，模型并行指的就是把模型的不同layer放在不同的機(jī)器上，不涉及將同一個(gè)layer切分到不同機(jī)器上的場(chǎng)景。

下圖是模型并行化，顯示了一個(gè)計(jì)算時(shí)間線，該示例有四臺(tái)機(jī)器和一個(gè)管道。

在正向階段，每個(gè)階段對(duì)本階段中的層的minibatch執(zhí)行正向傳遞，并將結(jié)果發(fā)送到下一階段。輸出級(jí)在完成前向傳遞后，計(jì)算minibatch的損失。
在后向階段，每個(gè)階段形成后向通道，逐一將損失傳播到前一階段。

worker 之間只能同時(shí)處理一個(gè) minibatch，系統(tǒng)中只有一個(gè)minibatch是活動(dòng)的，這極大限制了硬件利用率。

而且，模型并行化需要程序員決定怎樣按照給定的硬件資源來(lái)分割特定的模型，這其實(shí)無(wú)形之中加重了程序員的負(fù)擔(dān)。

1.2.3 Gpipe

除了通用問(wèn)題之外，GPipe 流水并行策略還有一個(gè)內(nèi)存問(wèn)題：需要緩存多份activation。

假如一個(gè)batch被分為 n 個(gè)micro-batch，則需要緩存 n 份activation。這個(gè) n 是梯度累加的次數(shù)，為了盡可能流水，通常這個(gè)累加次數(shù)都比較大，一般大于兩倍 stage 數(shù)目。那么即使只緩存少數(shù) Tensor，這種策略依然需要較多顯存。

0x02 論文

PipeDream 就針對(duì)這些問(wèn)題提出了改進(jìn)方法1F1B。PipeDream是第一個(gè)以自動(dòng)化和通用的方式將流水線并行，模型并行和數(shù)據(jù)并行結(jié)合起來(lái)的系統(tǒng)。PipeDream首先使用模型并行對(duì)DNN進(jìn)行劃分，并將每層的子集分配給每個(gè)worker。但是與傳統(tǒng)的模型并行不同，PipeDream對(duì)小批量數(shù)據(jù)進(jìn)行流水線處理，實(shí)現(xiàn)了潛在的管道并行設(shè)計(jì)。在任何時(shí)刻，不同的worker處理不同的輸入，從而保證了流水線的滿負(fù)荷以及并行BSP。

微軟在論文 PipeDream: Fast and Efficient Pipeline Parallel DNN Training 之中對(duì)于PipeDream進(jìn)行了詳細(xì)闡述，所以我們就基于此論文進(jìn)行分析。

2.1 方案概述

2.1.1 并行方式

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

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

每個(gè)階段對(duì)該階段中的所有層執(zhí)行向前和向后傳遞。PipeDream將包含輸入層的階段稱為輸入階段，將包含輸出層的階段稱為輸出階段。但是每個(gè)stage可能有不同的replication，這就是數(shù)據(jù)并行。

對(duì)于使用數(shù)據(jù)并行的stage，采用 round-robin的方式將任務(wù)分配在各個(gè)設(shè)備上，需要保證一個(gè)batch的數(shù)據(jù)在前向和后向發(fā)生在同一臺(tái)機(jī)器上，

2.1.2 1F1B

由于前向計(jì)算的 activation 需要等到對(duì)應(yīng)的后向計(jì)算完成后才能釋放（無(wú)論有沒(méi)有使用 Checkpointing 技術(shù)），因此在流水并行下，如果想盡可能節(jié)省緩存 activation 的份數(shù)，就要盡量縮短每份 activation 保存的時(shí)間，也就是讓每份 activation 都盡可能早的釋放，所以要讓每個(gè) micro-batch 的數(shù)據(jù)盡可能早的完成后向計(jì)算，因此需要把后向計(jì)算的優(yōu)先級(jí)提高，讓 micro-batch 標(biāo)號(hào)小的后向比 micro-batch 標(biāo)號(hào)大的前向先做。因此，如果我們讓最后一個(gè) stage 在做完一次 micro-batch 的前向后，立馬就做本 micro-batch 的后向，那么我們就能讓其他的 stage 盡可能早的開(kāi)始后向計(jì)算，這就是 1F1B 策略。

1F1B（one-forward-one-backward）的調(diào)度模式會(huì)在每臺(tái)worker機(jī)器上交替進(jìn)行小批次數(shù)據(jù)的前向后向計(jì)算，同時(shí)確保這些小批量在"后向傳播"時(shí)可以路由到"前向傳播"的相同worker。

這種方案可以使得每個(gè)GPU上都會(huì)有一個(gè)batch的數(shù)據(jù)正在被處理，使所有worker保持忙碌，而不會(huì)出現(xiàn)管道暫停，整個(gè)pipeline是比較均衡的。同時(shí)能確保以固定周期執(zhí)行每個(gè)stage上的參數(shù)更新，也有助于防止同時(shí)處理過(guò)多小批量并確保模型收斂。

0x03 流水線

PipeDream的pipeline parallelism（PP）是一種新的并行化策略，它將批內(nèi)并行與批間并行結(jié)合起來(lái)。

3.1 流水線改進(jìn)

我們首先看看流水線對(duì)于模型并行的改進(jìn)。

在上節(jié)的示例中，如果只有一個(gè)活動(dòng)的minibatch，系統(tǒng)在任何給定的時(shí)間點(diǎn)最多有一個(gè)GPU處于活動(dòng)狀態(tài)。

但是我們希望所有GPU都處于活動(dòng)狀態(tài)。因此PipeDream 多個(gè)小批量逐一注入到流水線中，從而通過(guò)流水線來(lái)增強(qiáng)模型并行訓(xùn)練。在完成小批量的前向傳遞時(shí)，每個(gè)階段都會(huì)異步地將輸出激活發(fā)送到下一階段，同時(shí)開(kāi)始處理另一個(gè)小批量。類似地，在完成一個(gè)小批量的向后傳遞后，每個(gè)階段都會(huì)將梯度異步發(fā)送到前一階段，同時(shí)開(kāi)始計(jì)算另一個(gè)小批量。

與普通層間并行訓(xùn)練相比，流水線有兩個(gè)主要優(yōu)點(diǎn)：

流水線通信量較少。流水線并行比數(shù)據(jù)并行的通信量要少得多。與數(shù)據(jù)并行方法（使用集體通信或參數(shù)服務(wù)器）中的做法（把所有參數(shù)聚合梯度并且將結(jié)果發(fā)送給所有worker）不同，流水線并行中的每個(gè)worker只需要在兩個(gè) stage 邊界之間將梯度和輸出激活的一個(gè)子集發(fā)送給另一個(gè)worker，這可能導(dǎo)致某些模型的通信量大幅減少。
流水線重疊了計(jì)算和通信。跨階段前向輸出激活和后向梯度的異步通信可以使得這些通信與后續(xù)小批量計(jì)算在時(shí)間上重疊，因?yàn)樗鼈冊(cè)诓煌妮斎肷线\(yùn)行，計(jì)算和通信完全獨(dú)立，沒(méi)有依賴邊，所以可以更容易的并行化。在穩(wěn)定理想狀態(tài)下，所有的 workers 時(shí)刻都在運(yùn)轉(zhuǎn)，不像模型并行化訓(xùn)練中會(huì)有停下來(lái)等待的時(shí)候。

下圖是實(shí)施了 1F1B 的流水線。Machine 1先計(jì)算藍(lán)色 1，然后把藍(lán)色 1 發(fā)送給 Machine 2 繼續(xù)計(jì)算。Machine 1 接著計(jì)算藍(lán)色 2。Machine 1 和 Machine 2 之間只傳送模型的一個(gè)子集。計(jì)算和通訊可以并行。

3.2 挑戰(zhàn)

PipeDream的目標(biāo)是以最小化總體訓(xùn)練時(shí)間的方式將流水線并行，模型并行性和數(shù)據(jù)并行性結(jié)合起來(lái)。然而，要使這種方法對(duì)大型DNN模型有效，獲得流水線并行化訓(xùn)練的潛在收益，PipeDream 必須克服幾個(gè)主要挑戰(zhàn)：

如何高效劃分流水線。與處理器中的流水線一樣，需要將DNN高效正確地劃分為若干“階段”（層序列），每個(gè)階段部署在不同的worker上執(zhí)行。
- 模型特質(zhì)和硬件拓?fù)鋾?huì)降低效率，劃分應(yīng)該具體取決于模型體系結(jié)構(gòu)和硬件部署。不好的劃分（階段性的工作量大范圍傾斜）可能會(huì)導(dǎo)致worker長(zhǎng)時(shí)間閑置。所以需要依據(jù)一定原則（通信和資源利用率）來(lái)劃分，比如：彼此有通信的層應(yīng)該分配到相鄰的處理器；如果多個(gè)層操作同一數(shù)據(jù)結(jié)構(gòu)，它們應(yīng)該被分配到同一個(gè)處理器上，彼此獨(dú)立的層可以映射到不同處理器上。所以分配算法也必須考慮模型特質(zhì)和硬件拓?fù)?/u>。

如何防止流水線瓶頸。

由木桶原理我們可以知道，在穩(wěn)定狀態(tài)下，一個(gè)流水線管道的吞吐量由這個(gè)流水線上最慢環(huán)節(jié)的吞吐量決定。如果各個(gè)環(huán)節(jié)的處理能力彼此差別很大，會(huì)導(dǎo)致管道中出現(xiàn)空閑時(shí)間（一半稱之為bubble），這樣最快環(huán)節(jié)必須停下來(lái)等待其他環(huán)境，會(huì)造成饑餓現(xiàn)象，從而導(dǎo)致資源利用不足。所以需要確保流水線中所有階段都大致花費(fèi)相同的計(jì)算時(shí)間，否則最慢的階段將會(huì)成為整個(gè)流水線的瓶頸。

如何在不同的輸入數(shù)據(jù)之間調(diào)度工作以均衡流水線。

與傳統(tǒng)的單向流水線管道不同，DNN訓(xùn)練是雙向的：前向傳播和后向傳播，兩種傳播以相反順序穿過(guò)相同層。如何協(xié)調(diào)流水線工作是一個(gè)問(wèn)題。

面對(duì)流水線帶來(lái)的異步性，如何確保訓(xùn)練有效。

流水線帶來(lái)的一個(gè)問(wèn)題就是weight版本眾多。在后向傳播時(shí)候如果使用比前向傳播時(shí)更高版本的weight來(lái)計(jì)算，則會(huì)造成訓(xùn)練模型質(zhì)量降低。
PipeDream管理后向通道里的權(quán)重版本，通過(guò)為每個(gè)小批量的weight維護(hù)版本號(hào)來(lái)解決這個(gè)問(wèn)題，這樣在后向通道里使用的權(quán)重版本就和前向通道里使用的相同，從而在數(shù)值上能夠正確計(jì)算梯度（我們后續(xù)文章會(huì)講解）。

3.4 流水線劃分算法

PipeDream基于一個(gè)短期運(yùn)行分析結(jié)果來(lái)自動(dòng)劃分DNN的層，依據(jù)分析結(jié)果使用算法來(lái)對(duì)不同階段之間的計(jì)算負(fù)載進(jìn)行平衡，同時(shí)最小化通信。PipeDream的自動(dòng)劃分算法總體目標(biāo)是輸出一個(gè)平衡的管道，確保每個(gè)階段大致執(zhí)行相同的總工作量。同時(shí)還必須確保各階段之間通信的數(shù)據(jù)量盡可能小，以避免通信中斷。算法如下：

將DNN層劃分為多個(gè)階段，以便每個(gè)階段以大致相同的速率完成，即花費(fèi)大致相同的計(jì)算時(shí)間。
嘗試以拓?fù)涓兄姆绞奖M量減少worker之間的通信（例如，如果可能，向更高帶寬的鏈路發(fā)送較大的輸出）。
因?yàn)镈NN并不總可以在可用的workers做平均分配，為了進(jìn)一步改進(jìn)負(fù)載平衡，PipeDream允許復(fù)制一個(gè)stage，即在這個(gè)stage上使用多個(gè)worker進(jìn)行數(shù)據(jù)并行。這樣多個(gè)worker可以分配到流水線同一階段，并行處理一個(gè)batch的不同的mini-batch，提高處理效率。因?yàn)閿?shù)據(jù)并行采用了RR，所以這套策略也被稱為 1F1B-RR(one-forward-noe-backward-round-robin)。

這個(gè)劃分問(wèn)題等價(jià)于最小化流水線的最慢階段所花費(fèi)的時(shí)間，并且具有最優(yōu)子問(wèn)題屬性：在給定worker工作量前提下，吞吐量最大化的流水線由一系列子流水線構(gòu)成，其中每一個(gè)子流水線針對(duì)較小worker工作量來(lái)最大化自己的輸出。因此，PipeDream使用動(dòng)態(tài)規(guī)劃來(lái)尋找最優(yōu)解。

具體如下圖：

3.5 Profile

DNN訓(xùn)練有一個(gè)特點(diǎn)：不同輸入的計(jì)算時(shí)間幾乎沒(méi)有變化。于是 PipeDream充分利用了這一事實(shí)，給定一個(gè)具有N層和M臺(tái)可用機(jī)器的DNN，PipeDream首先在一臺(tái)機(jī)器上分析模型，記錄向前和向后過(guò)程所花費(fèi)的計(jì)算時(shí)間，層輸出的大小以及每個(gè)層的相關(guān)參數(shù)的大小，最后輸出為一個(gè)結(jié)果文件。

分區(qū)算法不但使用profile結(jié)果文件作為輸入，而且還考慮了其他限制，如硬件拓?fù)浜蛶挕⒐と藬?shù)量和計(jì)算設(shè)備的內(nèi)存容量，最終將層分為多個(gè)階段，同時(shí)還確定每個(gè)階段的復(fù)制因子，以最小化模型的總訓(xùn)練時(shí)間。

所以總體算法大致如下：

因?yàn)镻ipeDream借鑒了很多GPipe的思路，所以可以看到其比Gpipe的進(jìn)步之處。

比如 Gpipe是通過(guò)在代碼中硬性預(yù)估ops來(lái)進(jìn)行流水線負(fù)載均衡，PipeDream則是先做profile，根據(jù)實(shí)際情況再做推理。

0x04 Profile階段

Profile是PipeDream工作的第一個(gè)階段，是分區(qū)算法的基礎(chǔ)。PipeDream根據(jù)profiling的結(jié)果，使用動(dòng)態(tài)規(guī)劃對(duì)模型進(jìn)行劃分，將模型劃分為不同的stage，以及每個(gè)stage的replication數(shù)。

這是PipeDream針對(duì)GPipe的一個(gè)改進(jìn)，兩者都是對(duì)每層的運(yùn)行時(shí)間進(jìn)行預(yù)估，然后對(duì)模型進(jìn)行劃分。

GPipe是利用經(jīng)驗(yàn)值或者數(shù)學(xué)的方法來(lái)對(duì)運(yùn)行時(shí)間進(jìn)行預(yù)估。
PipeDream根據(jù)profiling的結(jié)果對(duì)運(yùn)行時(shí)間進(jìn)行預(yù)估。

因?yàn)橛袑?shí)際數(shù)據(jù)進(jìn)行支撐，所以PipeDream更加準(zhǔn)確和先進(jìn)。

4.1 思路

評(píng)測(cè)機(jī)制利用了這樣一個(gè)事實(shí)：DNN訓(xùn)練在計(jì)算和通信時(shí)間上幾乎沒(méi)有變化。所以我們可以通過(guò)小批量數(shù)據(jù)的profile推理出DNN訓(xùn)練時(shí)間。為了確定所有層的運(yùn)行時(shí)間，PipeDream在其中一臺(tái)機(jī)器上使用1000個(gè)小批量對(duì)DNN模型的短期（幾分鐘）運(yùn)行進(jìn)行 profile。

4.1.1 如何計(jì)算

運(yùn)行時(shí)間

對(duì)于每一層的運(yùn)行時(shí)間，我們可以通過(guò) 運(yùn)行時(shí)間 = 計(jì)算時(shí)間 + 通信時(shí)間來(lái)得到。

計(jì)算時(shí)間就是每層layer前向和后向的計(jì)算時(shí)間，這個(gè)可以從profile得出。
通信時(shí)間就需要根據(jù)模型大小進(jìn)行估算，PipeDream 估計(jì)通信所需的時(shí)間為"需要傳輸?shù)臄?shù)據(jù)量"除以"通信鏈路上的帶寬"。

通信時(shí)間

在流水線上，大多數(shù)通信都有三個(gè)步驟：

1）在發(fā)送端機(jī)器上，從GPU傳輸?shù)紺PU移動(dòng)數(shù)據(jù)。

2）通過(guò)網(wǎng)絡(luò)從發(fā)送者到接收者發(fā)送數(shù)據(jù)。

3）在接收端，從CPU到GPU移動(dòng)數(shù)據(jù)。

而通過(guò)網(wǎng)絡(luò)從發(fā)送者到接收者發(fā)送數(shù)據(jù) 是最耗時(shí)的，所以PipeDream主要考慮這個(gè)因素。如果再對(duì)這個(gè)因素細(xì)分，則有：

對(duì)于從層 i 到層 i + 1 傳輸激活值的時(shí)間，PipeDream 基于 "激活值"來(lái)估計(jì)。
假如配置成了數(shù)據(jù)并行（對(duì)于層 i 使用 m 個(gè) worker 做數(shù)據(jù)并行）的情況，做權(quán)重同步的時(shí)間使用"權(quán)重"來(lái)估計(jì)：
- 如果使用分布式參數(shù)服務(wù)器，則權(quán)重?cái)?shù)量被預(yù)估為 4 x ( m - 1 ) x | w i | / m。
- 如果使用 all_reduce，則每個(gè)worker給其他workers發(fā)送 ( m - 1 ) x | w i | / m 個(gè) bytes，也接受到同樣數(shù)量字節(jié)。

4.1.2 Profile內(nèi)容

綜上所述，PipeDream在profile之中，為每個(gè)層 i 記錄三個(gè)數(shù)量：

Ti，層 i 的在GPU上向前和向后計(jì)算時(shí)間之和，即每層layer前向和后向的計(jì)算時(shí)間；
ai，層 i 的輸出激活的大小（以及向后過(guò)程中輸入梯度的大小）以字節(jié)為單位，即每層layer的輸出的大小；
wi，層 i 的權(quán)重參數(shù)的大小（以字節(jié)為單位），即每層layer參數(shù)的大小；

4.2 代碼

不同模型或者說(shuō)不同領(lǐng)域有不同的profile文件。

我們以 profiler/translation/train.py 為入口進(jìn)行分析。

4.2.1 訓(xùn)練腳本

以下我們省略了無(wú)關(guān)代碼。

4.2.1.1 訓(xùn)練過(guò)程

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計(jì)算網(wǎng)絡(luò)的計(jì)算參數(shù)等信息
        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)
                # 訓(xùn)練模型
                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 計(jì)算參數(shù)

上節(jié)在訓(xùn)練腳本制作的時(shí)候，torchsummary 的作用是計(jì)算網(wǎng)絡(luò)的計(jì)算參數(shù)等信息，對(duì)于 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 創(chuàng)建圖

create_graph 的作用就是使用torchgraph.GraphCreator創(chuàng)建一個(gè)圖，這個(gè)圖就可以理解為模型內(nèi)部的DAG圖，每個(gè)節(jié)點(diǎn)記錄如下信息。

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."""
    # 創(chuàng)建圖
    graph_creator = torchgraph.GraphCreator(model, summary, module_whitelist
    # 構(gòu)建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()
    # 運(yùn)行以得到profile                                        
    model(src, src_length, tgt[:-1])
    graph_creator.unhook_modules()
    # 輸出profile結(jié)果                                        
    graph_creator.persist_graph(directory)

4.2.2 創(chuàng)建圖

創(chuàng)建圖基本是在 GraphCreator 內(nèi)完成。

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 設(shè)置wrapper

hook_modules 的作用是給模型的forward函數(shù)設(shè)置一個(gè)wrapper，并且遍歷為子模塊設(shè)置，這樣在模型運(yùn)行時(shí)候可以跟蹤模型之間的聯(lián)系。

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):
                # 如果已經(jīng)被包裝，則插入input
                input.append(wrapped_inputs_list[i].tensor)
            else:
                key = wrapped_inputs_list[i]
                if key in this_creator.inputs: # 如果是原始輸入，則不進(jìn)行包裝
                    wrapped_inputs_list[i] = this_creator.inputs[key]
                else:
                    j = len(this_creator.inputs)
                    # 如果沒(méi)有被wrap, 則構(gòu)建一個(gè)TensorWrapper進(jìn)行包裝
                    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)
        # 對(duì)結(jié)果進(jìn)行包裝
        wrapped_result = TensorWrapper(result, str(self), this_creator)
        
        # 把邊添加進(jìn)入圖
        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

    # 遍歷子模塊，遞歸設(shè)置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.
            # 遞歸設(shè)置wrapper
            self.hook_modules(sub_module)
    if root: # 對(duì)于root進(jìn)行處理
        this_creator.forward_original_methods.update({module: module.forward})
        module.forward = forward_wrapper_root.__get__(module, module.__class__)

4.2.2.2 TensorWrapper

TensorWrapper 就實(shí)現(xiàn)了wrapper功能，graph_creator.summary 就是之前torchsummary.summary得到的網(wǎng)絡(luò)等信息。可以看到此類會(huì)遍歷 summary，計(jì)算 forward_compute_time 等信息，最終構(gòu)建了一個(gè) node。

需要注意的是：activation_sizes 是根據(jù) output_shape 來(lái)計(jì)算的。

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

對(duì)于某些內(nèi)置方法，則也會(huì)相應(yīng)處理，比如如下。

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

最終對(duì)應(yīng)：

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結(jié)果輸出到文件。

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)

具體調(diào)用了 graph.py 的函數(shù)完成，這里摘錄 to_dot函數(shù)如下：

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 結(jié)果

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

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階段的內(nèi)容，就是：運(yùn)行訓(xùn)練腳本，依據(jù)運(yùn)行結(jié)果來(lái)計(jì)算參數(shù)，建立一個(gè)模型內(nèi)部的DAG圖，然后把參數(shù)和DAG圖持久化到文件之中，后續(xù)階段會(huì)使用這個(gè)文件的內(nèi)容。

下一篇我們分析如何計(jì)算自動(dòng)分區(qū)。