本文分享自華為雲社群《深度解讀昇騰CANN記憶體複用技術，降低網路記憶體佔用》，作者： 昇騰CANN。
隨著大模型的興起，神經網路規模不斷擴大，對記憶體資源的消耗也越來越高，如何降低AI演算法的記憶體佔用成為AI計算領域降成本提效能的一個重要手段。當前降低AI演算法記憶體佔用的主要技術有記憶體複用（Memory Reuse）、重算（Recompute）、記憶體交換（Memory Swap）、控制執行序（Topo最佳化）等，其中記憶體複用是應用最廣泛的記憶體最佳化技術，該技術可在不影響網路執行效能的前提下降低記憶體佔用，結合Topo最佳化可以進一步降低記憶體佔用，且對執行效能影響也較小。
靜態shape的神經網路在基於計算圖模式執行時，由於具備計算圖的全域性視角，可以較完整地記錄和管理全域性記憶體及生命週期資訊，GE（Graph Engine）圖引擎基於業界常規的記憶體最佳化技術，藉助全圖視角對記憶體複用演算法進行了更好的尋優處理，同時進行了Topo排序最佳化降低網路記憶體理論最小值，從而更有效地降低網路記憶體佔用。
1 記憶體複用技術
1.1 常規記憶體複用技術
如下圖所示有一個簡單的網路，經過Topo排序後每個節點都有一個唯一的ID，每個節點的輸出記憶體都可以用這個ID來表達生命週期。
其中“生命週期”是指記憶體從分配到釋放的時間段，用拓撲排序序號來表示，在複用處理過程中，生命週期沒有重疊的才能複用。
“記憶體複用”是指按照生命週期和記憶體大小，把不衝突的記憶體重複使用，來降低網路記憶體佔用。記憶體複用演算法主要有Binary block複用和Max block複用兩種，兩種複用演算法原理類似，唯一差別是block的大小規劃有差異。這裡的“block”是指記憶體複用時的虛擬記憶體塊，並不是真正的記憶體。
記憶體佔用的“理論最小值”計算方式為：按照記憶體使用順序，在每個節點生命週期開始加上記憶體大小，生命週期結束減去記憶體大小，不考慮記憶體碎片和空洞等浪費問題，理論計算得到的記憶體佔用最小值。如下圖所示，記憶體佔用的理論最小值為4608，但是如果不做記憶體複用，總記憶體大小為8704。
原始網路記憶體需求示例圖：
記憶體複用處理的主要步驟：
1、獲取記憶體複用block的Range值
Max block演算法取所有記憶體大小的最大值作為Range，此例為[4096]，Binary block演算法會按規則生成多個Range，此例為[1024,2048,4096]。
2、Reuse處理
按照拓撲排序進行記憶體分配和複用，Max block複用演算法block大小都一樣，只要生命週期滿足要求就可以複用，Binary block複用演算法block大小不同，要生命週期和block大小同時滿足要求才可以複用。
3、Resize處理
複用處理時block大小一般都比實際分配的記憶體大，resize的作用就是根據複用結果把block大小改成合適的大小並做512位元組對齊。
下面以Max block演算法為例介紹原始記憶體複用技術的實現原理，Range為[4096]，因此所有block大小都為4096。Max block演算法記憶體複用示例圖如下所示：
Reuse處理：按照網路節點順序進行記憶體複用。記憶體分配時的block大小都是4096，標號為A、C、E的記憶體可以複用,而標號是B、D的記憶體由於生命週期存在交叉，無法複用。所以複用後得到3個block。
Resize處理：按照block中實際節點記憶體大小進行壓縮。複用完成後每個block取所有節點的記憶體最大值作為該block最終記憶體大小，如block1最大的是E，最終大小為4096，block2只有一個B最終大小為2048，block3為512，最後所有block的大小累加就是複用後的總記憶體大小。如下圖所示，複用後總記憶體大小為6656。
1.2 GE記憶體複用最佳化技術
Max block演算法進行記憶體複用時，要求所有block大小都是block Range，即網路節點的記憶體最大值4096，所以即使實際僅需要較小的記憶體也要按照最大值進行記憶體分配，存在一定的記憶體浪費。Binary block演算法雖然對block大小進行了細化，但是要在一定大小範圍內的節點才會複用，記憶體較大和記憶體較小都無法複用，也會影響記憶體複用的效果，存在記憶體浪費的弊端。
因此，GE採用了一種更極致的記憶體空間複用最佳化技術。一級複用時只把相同大小的記憶體複用在一起得到一組block，二級複用時把block按照從大到小排序，順序取後面一個block，如果生命週期不重疊和當前block進行合併，否則不做處理，以此類推進行多級遞迴判斷，最後把未被合併的block大小進行累加就是最終記憶體大小，由於大block合併多個小block，因此可以進一步減少記憶體佔用。
一級複用處理：按照實際大小進行復用，不同大小的block不復用在一起，因此AC可以複用，AE大小不同無法複用，複用後示例如下圖所示。
一級複用處理後得到4個block（比原有節點數減少了一個，此例比較簡單，實際網路中block的數量可以比節點數減少一個數量級），然後根據網路節點生命週期獲取block的生命週期。block生命週期示例圖如下所示：
然後把block按照從大到小進行排序，排序後結果如下圖所示。
二級複用處理：把排序後的大block和小block依次進行合併處理，一個大block可以合併多個小block。
- block4生命週期為[5,6]，block2的生命週期為[2,4]，生命週期不重疊，可以和block2進行合併，起始地址對齊。
- block1生命週期為[1,4]，和block4生命週期不重疊，但是和block2重疊，因此擺放在block2的右側，起始地址為block2的結束地址。
- block3生命週期為[4,5]，和block4重疊，不能和block4合併，因此只能擺放在block4的右側。
- 最終得到複用後的記憶體大小為4608（4096+512），和理論最小值一致，沒有任何浪費。
block合併示例圖如下所示：
1.3 GE記憶體複用最佳化效果
GE的記憶體複用最佳化主要有以下優點：
- 提高記憶體複用率，減少記憶體佔用
解決不同大小記憶體之間的組合複用問題，可以把一塊大記憶體切分成多塊小記憶體進行二次複用，進一步提高記憶體複用率，減少記憶體佔用。
- 降低複用演算法耗時
先按節點記憶體實際大小和生命週期進行一級複用，減少二級複用時的block個數（可以減少一個數量級），有效減少二級複用演算法耗時。
從16個實際網路模型的測試結果看，採用GE記憶體複用最佳化技術可以更好地降低網路記憶體，單Stream大部分場景可逼近理論最小值，而處理耗時也較短，如下圖所示。
較好的記憶體複用效果主要得益於GE計算圖模式有全域性的記憶體大小和生命週期資訊，能夠讓複用演算法進行更好的尋優處理；而單運算元模式下，由於演算法無法預測記憶體大小，生命週期也沒有全域性資訊，會導致記憶體空洞和碎片問題，無法將記憶體複用做到極致，大部分場景記憶體佔用只能做到理論最小值的1.2倍以內。相同網路圖模式記憶體佔用會比單運算元模式更少，例如，相同的網路下，若單運算元模式需要15.14G，圖模式只需要13.8G，記憶體佔用可降低9%左右。
2 Topo最佳化技術
2.1 常規Topo技術
記憶體複用最佳化技術主要解決如何逼近理論最小值的問題（一個網路在確定的Topo排序下都有一個理論最小值），而改變Topo排序可以進一步降低理論最小值的下限。業界常用的Topo排序演算法有BFS，DFS，逆DFS，在不同的圖結構下也會有較大的差異，常規排序演算法沒有考慮記憶體相關因素。常規Topo技術示例如下圖所示：
如上網路結果所示，DFS選擇左側分支先排序時，如圖10中的“優先排序左側分支”，會導致B的記憶體生命週期較長[2,6]，無法和C，D記憶體複用。DFS選擇右側分支先排序時，如圖10中的“優先排序右側分支”，B的記憶體生命週期是[5,6]，可以和C，D記憶體進行復用，但是又會導致A記憶體的生命週期變長由[1,3]變為[1,5]，因此常規Topo排序演算法無法做出針對記憶體的最優選擇，最後複用後記憶體大小是一個不確定的結果。
2.2 GE Topo最佳化技術
GE針對Topo技術進行了最佳化，首先做常規Topo排序，然後對排序結果進行處理，結合記憶體策略把某些節點的排序重新調整，從而確保調整後記憶體大小是可以降低的。其中一個最佳化策略就是把輸入是長生命週期的節點向後移動，移動到排序最靠前的一個輸出節點前。
長生命週期記憶體的判斷規則：
- 固定規則，根據型別判斷，如果某節點的輸入是Variable、Constant、Const型別，或者不做複用的Data、RefData型別，則此節點可向後移動，因為這些型別的節點記憶體本身不做複用，所以節點向後排不影響記憶體生命週期。
- 動態規則，根據Topo排序結果進行判斷，排序完成後，理論上可以計算出某節點的輸入節點記憶體的最大生命週期，標記為L1，然後尋找本節點的輸出記憶體的最小生命週期，標記為L2，如果L1 > L2，則可以將本節點移動到L2前。
如下圖所示，假設節點A是Variable型別，則可以將節點B移動到輸出節點F的前面，進行排序最佳化。
2.3 GE Topo最佳化效果
做了Topo最佳化後，降低了網路的記憶體理論最小值，最後記憶體複用後的大小會進一步下降，如下是幾個網路的實際測試結果。
3 更多介紹
GE記憶體複用技術的相關介紹就到這裡，歡迎大家關注後續技術分享。如需獲取更多學習資源請登入昇騰社群。
點選關注，第一時間瞭解華為雲新鮮技術~