導讀 本文將介紹 Oneflow 基於重計算的動態圖視訊記憶體最佳化工作——Coop。Coop 的核心創新是實現動態圖重計算策略和視訊記憶體分配機制的聯合最佳化。
文章包括以下四大部分：
1. 動態圖重計算（DTR）介紹
2. 現有方法的侷限性
3. OneFlow Coop 的做法
4. 實驗效果展示
分享嘉賓｜張建浩 一流科技 框架開發工程師
編輯整理｜孫彪 北京航空航天大學
出品社群｜DataFun
Coop，是對於動態圖重計算策略，還有視訊記憶體分配策略的聯合最佳化。動態圖重計算，簡稱 DTR。首先介紹一下動態圖重計算的背景。
在大部分場景下，神經網路訓練過程中佔用的視訊記憶體量，大部分都來自於網路訓練的中間特徵，而不是網路的引數（即權重）。例如，ResNet-50 的權重總大小就只有不到 100MB，剩下的視訊記憶體其實都是由訓練過程中在前向時產生的那些中間特徵佔據的。這些中間特徵之所以會留在視訊記憶體裡面而沒有被及時釋放掉，是因為在反向傳播的過程中還會再用到它們。可以看右邊的示例圖，這個是來自於陳天奇的一篇奠基性的論文 Sublinear 裡的一張圖。它有兩列，左邊這一列是前向的過程，從 input 到最後的 loss。右邊這一列是反向的過程，從 label 一直到得到所有的梯度。可以看到，在右邊反向這一列的反向過程中，會用到前向過程中的一些中間結果，這也就是為什麼這些中間特徵沒有被釋放，是因為其實後面它們還會被用到，所以暫時留在視訊記憶體裡面。在動態圖的場景下，對於這些中間特徵，我們可以先釋放掉一部分的中間變數。然後在後向過程中，等到要用到這些被釋放掉的中間變數的時候，再根據前向記錄下來的計算過程，把中間變數重新計算出來。這樣會付出一定的額外的計算開銷，但是減少了在訓練過程中所需的視訊記憶體佔用。
動態圖重計算是由這一篇 paper 開創的。其思路就類似於 cache。我們常見的 cache 會有一個預先設定好的容量，即 cache 容量。如果 cache 的容量已滿，還要再往 cache 裡面再插入一個新的元素，就會從這個 cache 已有的元素中挑出來一個，把它刪除掉，然後再把新的元素插進去。至於是選 cache 裡面的哪一個元素來釋放，就可以用一些非常簡單但是卻很有效的策略，比如常用的 LRU 策略。它的選擇策略很簡單，就是根據最近一次被訪問的時間，即距離現在最近的一次訪問過了多久。如果過的時間越長，就認為它是最不應該被 cache 住的，所以會把它給釋放掉。神經網路框架的動態圖重計算也是類似的思路，它也是用一個非常簡單的判斷標準：當視訊記憶體已經滿了的時候，會用一個非常簡單的判斷標準去找出一個我認為釋放代價最低的一個 tensor，把它釋放掉。釋放掉之後，再去試一試，看空閒視訊記憶體是不是已經夠了，如果已經夠了就結束了，如果還是不夠，就會去迴圈地做上述這個過程，直到足夠的空閒視訊記憶體被釋放出來，使得後續的視訊記憶體申請能夠成功。DTR 的判斷標準考慮了三個因素：第一個是重計算的開銷越小越好，這個好指的是 tensor 的釋放代價低，代價越低，它就越應該被釋放；第二個是佔用的視訊記憶體越大越好；第三個是上一次訪問距離現在的時間越長越好。這 3 個都是很符合直覺的考慮因素。DTR 會綜合考慮這 3 個因素去決定哪一個 tensor 釋放代價是最低的。其中還有一點是重計算的開銷，DTR 為了防止出現很長的一條 tensor 鏈被連續地釋放掉的情況，除了考慮 tensor 自己的重計算開銷之外，還會遞迴地將不在視訊記憶體內的父代 tensor 的計算開銷考慮在內。可以看右邊這張圖，這個圖的 T0 到 T6 就形成了一個計算圖。現在想計算出 T7，可以從 T5 和 T6 算出 T7。這時候如果發現視訊記憶體不夠了，它就會從灰色的，也就是目前在視訊記憶體裡的那些 tensor 裡面，去挑一部分釋放掉，為 T7 騰出空間。這是一個在執行時去動態做的過程。例如，在這張示意圖裡面，它就挑了 T2 和 T3 來釋放，為 T7 騰出了空間。所以算完 T7 之後，它的狀態就變成了這些，灰色的這四個在視訊記憶體裡面，T2 和 T3 就已經被踢出去了。
在我們的實踐過程中，我們發現已有的工作即 DTR 是有一些侷限性的。首先，它是一種帶有盲目性的貪心演算法，因為每次只會考慮當前最該被釋放即代價最低的那一個 tensor，再不斷地去做迴圈直到視訊記憶體申請能夠成功。但是迴圈的每一步之間，每一步被釋放的這些 tensor 之間的關係卻沒有被考慮進去，就會經常出現這樣的場景：比如需要騰出一個 100MB 的空間，開始迴圈，先挑了一個 60 MB 的 tensor 釋放掉，又挑了一個 40MB 的 tensor 釋放掉，這時雖然看起來 40 加 60 等於 100，能夠容納 100MB 的 tensor，但是實際上卻不是這麼簡單。因為在視訊記憶體裡面我們對視訊記憶體的連續性也是有要求的。我們單一地分別釋放一個 60MB 的和 40MB 的 tensor，很可能它們釋放出來的空閒視訊記憶體是兩段分散的。這樣分散的視訊記憶體我們是用不了的，也就是常見的視訊記憶體碎片，而不是能夠容納 100MB tensor 的一段連續的大視訊記憶體。既然 DTR 用不了這兩個碎片，那麼這個演算法還會繼續去迴圈，去找其他的 tensor 再釋放，直到正好釋放了一整個 100MB 的 tensor，或者兩個或多個小的 tensor 正好一起湊夠了 100MB 的連續視訊記憶體，這顯然是一個不太好的性質。
第二個侷限性就是它的重計算演算法和 tensor 在視訊記憶體池裡面的排列方式是有一定衝突的。一般來說，假設網路裡面有一個直線型 OP 序列，直線型 OP 序列產生的 tensor a，b，c，d 被放在視訊記憶體池裡。對於視訊記憶體池裡的一個 tensor，它被釋放的時候，它的視訊記憶體不會直接歸還給系統，還會在這個視訊記憶體池裡面，視訊記憶體池會把它給 hold 住。後續如果再有人要新的視訊記憶體，就直接從視訊記憶體池裡 hold 住的那些視訊記憶體裡面去直接分配，就不走作業系統了。也就是說，顯著池會存在一個很大的空閒視訊記憶體塊。這樣這些順序的 tensor a，b，c，d 就會在很大的空閒視訊記憶體空間裡面連著排。所以如果 tensor 在網路裡面是連續的，那麼它在視訊記憶體池裡面即物理上的記憶體排列按照預設規則也會是連續的。這樣會引起的問題可以參考下圖。
（a）圖是一個比較典型的網路結構，即 relu-conv，這裡就省略了 BN，因為畫起來要簡單一些，但是不影響本質的分析。一般來說，它在視訊記憶體池裡面的排列也會是按照這個順序 relu-conv-relu-conv…。此時如果新來一個 X5，它的視訊記憶體大小是 2 倍的小方塊的長度，我就要挑兩個 tensor 釋放。如果按照經典的 DTR 的做法，它會把 X0 和 X2 這兩個 tensor 先釋放掉。這個背後有兩個原因。一個是比較顯而易見的，因為 X0 和 X2 它們都是 relu 產生的，relu 的重計算代價比 conv 的重計算代價要低很多。另一個原因是我先挑了 X0 來釋放，前文講到，計算重計算代價的時候，裡面的 c，就是重計算的開銷，它不僅是 tensor 本身的計算開銷，也會包括不在視訊記憶體內的那些父代 tensor 的重計算開銷。也就是說當我釋放了 X0 之後，它的子代 tensor X1 的重計算代價也會增加，使得 X1 更加難以被釋放了。能看出來這其實是一件很矛盾的事情，我們想要的是 X0 和 X1 都被釋放掉，這樣才能夠形成一段連續的空閒視訊記憶體。但是 DTR 對重計算代價的設計和 tensor 在視訊記憶體池裡的排列方式卻一起作用，使得我們反而更加難以去形成連續的視訊記憶體，這就是它的第二個侷限性。它帶來的結果就是，可能先釋放掉 X0，X2，接下來還得再釋放一個，才能夠形成一段連續的視訊記憶體容納 X5。如果運氣好，可能下一次就選到了 X1，但是也有可能又選到了 X4 或者後面別的 tensor，那就更糟糕了，因為可能要釋放很多無辜的 tensor，才能夠碰巧形成一段連續的空閒視訊記憶體。
第三個侷限性是動態圖重計算的傳統方法 DTR 裡面，對 in-place op 也是有一定影響的。因為在重計算的時候，相當於是在重新計算歷史上的 op。in-place op 會就地修改掉一個已有的 tensor 的值，而不是產生新的 tensor。這樣一來，在重算的時候，拿到的資料可能就是一個已經被修改過的資料，而不是當時所用的資料了，即現在的資料可能已經被 in-place op 給修改過了。這樣它們就會有一個相互的作用。參見上圖，現在有這樣的一個計算軌跡，從 p 到 x 到 y，分別減 2 加 1。現在如果來了一個 in-place op，對 x 做 in-place relu。本來計算的軌跡是減 2 加 1，比如 y 被釋放了，如果被釋放了，要重計算 y，從 x 加 1 就能把 y 再重計算出來。但如果對 x 做了 in-place relu 之後，x 的值就變了，這個時候再對它加 1，得到的就不是我想要的 y 了，資料就錯了。在 DTR 的做法裡面，引入了一個 copy-on-write 的中間層，實質上把 in-place op 變成了一個非 in-place op，也就是 out-place op。它的具體做法是：申請一塊新的空間去放置一個 x’。如果對 x 做 in-place 操作，它會申請新的空間放置 x’。在 x 和 x’ 之間去做一個普通的 op，也就不 in-place 了。這樣一來，在計算圖或者在計算軌跡上面，x 和 x’ 就變成是分離的兩個沒有關係的 tensor，這樣就能保證它是可以被正確地重計算的。這樣，首先，它失去了 in-place op 的優點，比如節省視訊記憶體，對 cache 更加友好，這些優點全都沒有了；另外，在模型訓練過程中，模型的權重是不能被釋放的，釋放的代價實在是太高了，而模型權重通常都會 in-place 地被更新，如果我們也用這個 copy-on-write，也就是不 in-place 更新了，而是 out-place 更新，就會生成很多新的 tensor，這些新的 tensor 就會分散地分佈在視訊記憶體池裡面的各個地方，這樣就相當於把視訊記憶體池裡面的視訊記憶體給割裂了，因為這些 tensor 是不能被釋放的。這樣就更加難以在視訊記憶體池裡面形成一段很長的連續的空閒視訊記憶體了。以上就是 DTR 的侷限性。在我們的工作裡面，聯合考慮兩件事情，一個是視訊記憶體的排列方式，即 tensor 的視訊記憶體在視訊記憶體池裡的排列方式，以及重計算的問題，就能夠把上面的三個侷限性都解決掉。
首先，我們考慮了視訊記憶體的排列方式。不再是盲目地去迴圈，每次找出代價最低的一個 tensor 了，而是整體的考慮，當要釋放一個 tensor 的時候，它的周圍是怎樣的，相當於是在找出一個 tensor 的集合，釋放這個 tensor 集合之後，能夠形成足夠的連續視訊記憶體，而且這個集合的釋放代價是最低的，也就是上圖中的公式。後面的 subject to 是一個限定條件，這個 M(S, L) 限定條件是指在視訊記憶體排列方式 L 下，釋放掉集合 S，集合裡面的所有 tensor 能夠形成期望的連續空閒視訊記憶體的長度。在滿足限定條件的情況下，想找出最優的 S 和 L，這個最優即集合 S 裡面的所有 tensor 的釋放代價加起來最小。
基於上述目標，我們提出了三個模組，分別是 recomputable in-place，op-guided tensor allocation 和 layout-aware eviction。上圖是視訊記憶體池的狀態，即記憶體的排列狀態。顯示池裡面綠色的是釋放代價低的 tensor，藍色的是釋放代價高的 tensor，紅色的是不可被釋放的 tensor，這些紅色的 tensor 主要是 parameters、buffers 等。可以看到最開始的時候，這 3 類 tensor 是交叉隨機排列的。經過 recomputable in-place 之後，這些不可釋放的 tensor，會一直留在視訊記憶體池最邊緣的地方，這樣就能夠促使我們形成足夠的連續視訊記憶體。再接下來，經過 op-guided tensor allocation，就能夠把 cheap 的 tensor 還有 expensive 的 tensor 給分散到兩邊。這樣就很容易找出釋放代價低的 tensor 也就是 cheap 的這部分 tensor 了。最佳化 tensor 的排列方式，如果對應到前面的公式，就是調整公式裡面的L，再透過 layout-aware eviction，也就是透過滑動視窗，就可以找出想要的一個連續 tensor 的集合。
對於 recomputable in-place，首先它是可重計算的，與 naïve 計算方式不同。並且，它又是 in-place 的，具有 in-place op 的那些優勢，比如節省視訊記憶體和對 cache 友好。具體做法是受啟發於一個 PL 領域的工作，一個叫 Koka 的程式語言，其中引入了一個 functional but in-place 的機制，類似於尾遞迴最佳化。這是一種通用的最佳化方式，具體來說，假如已經知道某一個變數接下來不會被用到了，就可以把新的變數直接分配到已經知道不會被用到的變數的記憶體地址上面，這就相當於對已經確定未來不會被用到的變數的記憶體做一個複用，這樣就達到了類似於 in-place 的效果。但是如果單看變數，前後確實是不同的兩個變數，但這兩個變數共享了一個記憶體空間，因為他們的生命週期不重疊。Coop 裡面的 recomputable in-place 正是受此啟發，不過這裡不是變數，而是 tensor，但核心思想類似，讓兩個 tensor 暫時共享同一個記憶體空間。現在在這個 recomputable in -place 裡面，雖然也是有兩個不同的 tensor X 和 X’，在這兩個 tensor 之間，表面上是在做一個非 in-place 的操作，去形成一個正確的計算軌跡。但是實際上，因為這兩個 tensor 共享了同一個記憶體空間，所以效果是和 in-place OP 完全一樣的。等這個 in-place OP 結束了之後，再和原來 DTR 的 copy-on-write 一樣，把輸入 tensor 設定為已釋放的狀態，最終達到的狀態和 DTR 的 copy-on-write 完全一樣。但是我們仍然保留了它 in-place OP 的良好的性質，也就避免了前面所提到的 copy-on-write 的 2 個缺點。
接下來是 op-guided Allocation。我們根據產生 tensor 的 op 的性質去指導我們把 tensor 放在視訊記憶體池裡面的什麼地方。我們發現如果計算 cost density，即代價的密度，也就相當於是單位視訊記憶體下重計算的開銷。可以看到 Conv 和 MatMul 跟其他的 op 相比，代價密度有一個數量級上的差別。這就引導我們把 Conv 或 MatMul 生成的 tensor 和其他 op 生成的 tensor 劃分成兩類，分別放在不同的位置。在我們的實現裡，我們利用了重計算的一個特性，即最大視訊記憶體閾值是預先給定的，這樣我們就可以從左右兩端去放 tensor，而不是隻能一直放在左邊。具體來說，比如 x0 它是 ReLU 產生的，我們把它放在左邊，而 x1 是 Conv 產生的，放在右邊，以此類推。這樣代價密度小的 tensor 在一邊，代價密度大的 tensor 在另一邊，我們就可以更加容易地找到代價之和更低的 tensor 集合。同時它也避免了前文提到的自相矛盾的問題，即 x0 的釋放反而增加了 x1 的代價，導致 x0 和 x1 難以同時釋放形成連續的空閒視訊記憶體，我們的放在兩邊的方式就可以避免這個問題，因為現在 x0 和 x1 在視訊記憶體池裡面就已經不再連續了。雖然 x0 釋放以後，x1 的代價還會上升，但是已經不影響連續空閒視訊記憶體的形成了。
最後一個模組是 layout-aware eviction，滑動視窗。為了能夠做滑動視窗，我們做了一些預處理，把問題轉化了一下。具體來說，在視訊記憶體池裡面，我們把那些已有的空閒的視訊記憶體塊看作一個代價為 0 的特殊 tensor，這時候 tensor 就包括那些空閒的視訊記憶體了。把所有的 tensor 按照在視訊記憶體裡面的地址排列成一個 list，如果兩個 tensor 在 list 裡面是連續的，那麼在物理視訊記憶體上面也是相鄰的。問題就變成了如何在這個 list 裡面找一個連續的子序列，子序列要滿足視訊記憶體之和大於預先給定的閾值 MR，還希望它的代價之和最小。這其實是一個很經典的演算法題，有兩個指標，每一次都判斷視訊記憶體之和有沒有大於 MR，如果小於，就把尾指標往右移一格，如果大於，就把頭指標往右移一格。在移動的過程中，邊移動邊記錄當前的兩個指標之間的 tensor 的 ht 即視訊記憶體之和，同時記錄 ht 之和最小的一組 tensor。於是一次遍歷就能夠得到最優的 tensor 的集合，而不像傳統的 DTR 方法要去不斷地迴圈遍歷。從理論上可以計算出來，如果期望的視訊記憶體大小長度符合一個均勻分佈，DTR 方法的時間複雜度會是 O(n2)，這個 n 指的是視訊記憶體池裡面有多少個 tensor，但是我們的 Coop 方法的時間複雜度則只有 O(n)，即只需要線性的時間就能夠找到想要的集合。
接下來看一下我們的實驗效果。藍色的是 Coop，綠色的是 DTR，黃色的是 DTE，DTE 是 MegEngine 團隊對 DTR 做的一個改進。橫軸是視訊記憶體的比例，假設不重計算的時候視訊記憶體佔用量是多少，現在如果只想用不重計算時的80%的視訊記憶體，那麼橫軸就是 0.8。縱軸是計算時間的比例，以某一個視訊記憶體閾值重計算之後，計算時間擴大了多少倍，比如擴大 1.1 倍、1.2 倍，線越往左下角越好。可以看到 Coop 在 8 個網路裡面都比較明顯地超過了另外兩條線。
另外比較有趣的一點是，我們統計了視訊記憶體碎片率，即在視訊記憶體池裡面空閒視訊記憶體佔的比例。可以看到 Coop 有一個數量級的提升。BiLSTM 是對數座標，其他兩種方法視訊記憶體碎片率已經快要爆表了，但是 Coop 仍然是一個很低的狀態，只有 10% 左右。另外三個網路也是類似的效果，Coop 的視訊記憶體碎片率都是很低的。
前文提到，去找 tensor 集合的過程也是有時間開銷的。傳統方法的時間複雜度是 O(n2)，而 Coop 方法是 O(n)，在實驗中也能夠得到證實。在大部分情況下，Coop 去找 tensor 集合的時間即搜尋時間都會比其它兩種方法要好。
目前我們剛剛投了一篇 paper，對應的程式碼還沒有合併到 OneFlow 的 master 分支裡面去，不過預計很快。大家如果感興趣可以掃碼加入我們 OneFlow 的微信群或者 QQ 群。問答環節
Q1：OneFlow 現在有使用 DTR 的場景嗎？印象裡 OneFlow 是靜態的圖。A1：曾經 OneFlow 是靜態圖，現在我們 OneFlow 是動態圖靜態圖都有。個人覺得未來的趨勢動態圖仍然會是主導地位。就像 PyTorch 一直是 eager-first 的策略，也一直是統治性的框架。而我們目前的 API 也是追求和 PyTorch 百分之百完全相容，你可以直接 import oneflow as torch。相當於可以認為我們是 PyTorch 很好的一個增強版。A2：支援的，用我們的 nn.Graph 包一層就可以。因為預設的時候我們和 PyTorch 一樣是靜態圖，我們有自己的 nn.Graph，用 nn.Graph 把 nn.module 包一層，就可以把動態圖變成靜態圖。Q3：Coop 是用計算換視訊記憶體資源嗎？主要的應用場景是什麼？大模型訓練嗎？
A3：對，可以這麼說。大模型訓練是一個典型的場景。也有一些場景比如：可能有些個人使用者，他的顯示卡不是很好，可能視訊記憶體很少。但是那種常見的模型，比如檢測模型，可能需要很大的視訊記憶體量才能夠訓練起來。這種場景下雖然不是大模型，但是對於這種卡的視訊記憶體很少的一些普通使用者來說，也是會有很大的幫助。
來自 “ ITPUB部落格 ” ，連結：http://blog.itpub.net/70027824/viewspace-2945265/，如需轉載，請註明出處，否則將追究法律責任。