摘要：近來，增大模型規模成為了提升模型效能的主要手段。特別是NLP領域的自監督預訓練語言模型，規模越來越大，從GPT3的1750億引數，到Switch Transformer的16000億引數，又是一個數量級的增加。
本文分享自華為雲社群《一文帶你瞭解MindSpore支援的萬億級引數超大模型關鍵技術！》，原文作者：HWCloudAI 。
前言
近來，增大模型規模成為了提升模型效能的主要手段。特別是NLP領域的自監督預訓練語言模型，規模越來越大，從GPT3的1750億引數，到Switch Transformer的16000億引數，又是一個數量級的增加。
模型規模的數量級的增大，雖然取得了一定程度的效能提升，甚至產生了某些意想不到的“神奇”效果（如GPT3），但其背後的計算開銷成了最大的問題，比如GPT3訓練使用了萬級的GPU和數週的訓練時間。如何既能利用超大規模的引數來提升模型表達和效能，又能控制計算量的較小的增加，成為了最主要的挑戰之一。以MoE為代表的動態神經網路技術被重點引入。大腦是典型的低能耗高效率的計算模式，稀疏啟用是最重要的特性。除了巨型模型在訓練推理特別是訓練時的計算效能挑戰外，當前巨型模型的訓練優化演算法另一個更大的挑戰是（不在此處討論），BP演算法是當前最為可用的深度網路優化，但更理想的優化演算法需要高並行、優化過程非對稱、並能夠在時空維度通過區域性持續優化完成整體優化。
1. 傳統的神經網路模型，前饋的時候，輸入的batch中，每一個樣本的處理，都將啟用網路中的每一個引數參與計算。
2. 條件計算最寬鬆的定義，指僅啟用網路中某些部分的一類演算法。Conditional Computation refers to a class of algorithms that activate only some of the different parts in a network. 在具體某類條件計算實現中，條件選擇模式，可能按照輸入的batch中每sample獨立啟用網路不同部分，可能按照輸入資料空間上不同的部分（比如image不同區域或者channel），可能按照輸入資料時間上不同的部分（比如time series的不同slide window或者video的不同的frame。），可能按照目標任務的不同每task獨立的，可能按照非可學習的固定的隨機分配不同的子網獨立計算。
3. 對不同的輸入（原始或者前層），按照一定條件，選擇性的執行後續部分網路的計算，這個技術下，有一些近似或相關的技術，如：dynamic neural network(s), conditional computing, conditional activation, sparse activating, selective execution, mixture of experts (MoE), dynamic routing, …；強相關的一些模型比如 Switch Transformer等。
條件計算的分類（廣義）
1. 按照routing是否可學習可以分為：learnable routing conditional computation和 unlearnable routing conditional computation.
2. 按照activation是否不執行non-activation計算，可以分為：hard conditional computation和soft conditional computation。對於hard-mode的條件計算，通過tensor挑選切分等操作，無論何種條件選擇模式，不需要啟用的資料將完全不參與不啟用的網路部分的計算；soft-mode的條件計算，可能僅採取將相關資料置零等方式來避免產生計算效果，但還是和不需要啟用網路部分實際執行計算過程。
條件計算的主要優勢
1. 計算有效，降低能耗：通過部分啟用部分計算，以每樣本條件啟用的條件計算為例，單個樣本只需要經過整個SuperNet的一部分參與計算。
2. 更大網路，表達更強：由於一處到多處的Route，各處（層）的Input被路由到不同的子網獨立計算，不同的輸入的相互在各層的表達相對獨立沒有影響，表達能力更強，網路可以更大，但表達效率降低了。
條件計算的網路和計算形式
條件計算的網路和計算形式比較靈活，部分構建形式如：（此處省略具體模型和論文引用，參見: http://intellabs.github.io/dis）
1. 按照CV等task的特點，用多個獨立的CNN作為expert網路，按照task來獨立路由，尾部組合後給一個大網路。
2. 使用更復雜的cascading等形式組合不同層級的不同的expert網路。
3. 通過決策樹等方法做資料變換實現路由。
4. 通過可學習的網路來選擇路由。其中策略學習的損失有多種構建形式：直接使用分類等任務的主損失，對不同專家的重要性和負載構建損失作為輔助損失等等。
條件計算的路由策略
1. non-learnable/hard-mode，通過某種確定性策略，如LSH等方式計算路由。
2. learnable-mode，通過可學習網路計算路由。網路規模可大可小，簡單的可學習路由為單層權重：G(x) = P(X*W)，G(x)為路由Gate函式，X為輸入， W為通損失函式來度量的可學習路由權重，P為某種挑選函式（如topk, sort等），在實際實現中，X*W的輸入與權重計算結果可能作為後續網路的輸入資訊的一部分，不僅僅利用G(x)來選擇路由，則需要對X*W的結果做歸一化，更典型的形式則為：G(x)=P(N(X*W))，其中N為表達Normalization函式，如Softmax。
條件計算的冗餘策略
條件計算的冗餘策略，可分為無冗餘條件計算和冗餘條件計算：
1. 無冗餘條件計算可通過P(.)函式的實現如topk(k=1,…)來實現；
2. 冗餘條件計算，可以多種實現形式，可以通過P(.)函式的實現如topk(k=n,…)，n>=2來實現，也可以通過硬冗餘模式，整個網路中支援輸入的複製和多路計算實現。
條件計算的挑戰
1. 路由演算法對模型質量的影響無論輸入和路由權重作用的資訊（X*W），是僅作為路由選擇並作為後續網路單元的輸入，還是直接作為後續網路單元的輸入的一部分，路由演算法決定了輸入資訊的處理流向，對模型的整體質量都有很大影響。2. 路由(routing)/門(gate)的穩定性隨機初始化的路由/門的權重，權重自身在不斷被訓練調整；在前後層的網路持續訓練變化，同一樣本在訓練的不同階段會被分派到不同的後續網路單元中，這種動態變化過於劇烈，將嚴重影響整個網路訓練過程的穩定性和收斂速度。3、路由的專家樣本重要性和負載的平衡性
訓練階段，每專家和樣本批次中樣本的關聯度重要性，和每批次中樣本被均衡分派到不同專家的負載平衡性，這兩個指標既相關又衝突。需要分別構建損失函式作為輔助損失，來優化這兩個指標。在arxiv:1701.06538《Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer》做了相關討論。
關於條件計算/動態神經網路
關於條件計算/動態神經網路，更多的資訊在《Dynamic Neural Networks: A Survey》arxiv:2102.04906 (http://arxiv.org/abs/2102.0490)一文中，作者對廣義的動態神經網路，將各種動態網路相關的技術按照例項級、時間級、空間級做了分類。
- 1. Instance-wise Dynamic NN：逐例項動態，每樣本獨立啟用不同的網路和引數（MoE為這個方向）。Dynamic Architecture：Dynamic Depth、Dynamic Width、Dynamic Routing/MoE；Dynamic Parameter：Parameter Adjustment、Parameter Prediction、Dynamic Feature(s)
- 2. Spatial-wise Dynamic NN：空間級動態：影像等不同空間位置啟用後續不同網路和引數。(CNN等)：Pixel Level、Region Level、Resolution Level
- 3. Temporal-wise Dynamic NN：時間級動態：時序資料按時序維切分啟用後續不同網路和引數。（video-frames, text-sequence, time-series, stream, ...)Text-SequenceVideo-Frames
上述為該綜述論文對Dynamic NN的總體分類。
從超大規模網路動態網路技術支撐角度，高表達能力，低計算代價為主的來考慮分類，從兩個維度對動態網路技術分類:
1. 按照在前饋計算時是否部分啟用：
Hard-Dynamic：在前饋的時候，部分網路絕對不啟用參與計算
Soft-Dynamic：在前饋的時候，部分網路經過softmax等gate/route後，通過張量元素置零等方式，失去表達能力，但會參與計算。
2. 按照動態啟用判定演算法的輸入：
- 逐樣本級：（在輸入層）按照每樣本的例項來決定動態網路的後續啟用。
- 亞樣本級：（在輸入層）樣本內時間/空間級啟用不同的後續網路單元。一般深度網路，不僅在輸入層會被選擇性啟用執行，在中間層也類似。
其中，智慧平臺支援Hard-Dynamic逐樣本級的動態神經網路，能比較自然的獲得網路結構大顆粒的稀疏啟用，在超大模型中能實現訓練和推理的高能效。
動態神經網路相比與靜態結構的神經網路，在相關研究中，從效能，表達，泛化、魯棒，可解釋等方面做了大量對比研究。從智慧平臺通過計算成本儘量低的支援超大規模網路來提升模型效能的角度看，Efficiency和Representation最為重要：
1、Efficiency：靜態網路“牽一髮而動全身”，每一個樣本輸入整個網路/所有引數都要響應，這對超大網路來取得領先效果的模型能耗挑戰太大。
2、Representation: 引數量更大，表達容量更大；但MoE等結構在深度網路的各層特徵的表達上，複用降低，每引數的表達效率更低。
實現策略
實現各種模型的帶有動態路由稀疏啟用的超大規模引數版本，需要分模型研究和實現。
以Switch Transformer為例，其引數擴充套件到部分在Transformer的FFN部分。其MoE化擴充套件，如下圖：
(圖片來源：Switch Transformer論文)
可見，MoE化主要變化在需要Expert子網路前後增加MoE相關的邏輯。本文主要介紹平臺上的實現。動態路由條件計算，主要包括四個步驟：路由計算、資料分派、獨立計算，結果合併。
1. 路由計算-Gate：根據輸入（可以為整個網路的輸入，或者前面網路單元/層的輸出），在路由單元完成計算，在以batch內sample-wise的路由中，計算出每個樣本要分派的後續網路路由（Mixture-of-Experts/MoE中的專家）。
2. 資料分派-Dispatch：從輸入的整體的Tensor中，按照路由計算的樣本-專家關係，收集合並出每個專家需要處理的Tensor。如果在固定expert-batch的設計中，要平衡每批訓練中，分派到每個專家的樣本數和專家每輪訓練最大容量，由於樣本輸入的隨機性，很難保證較為均勻的分派，對於低於最大容量的批次，對固定batch-size的要做pad，對於高於最大容量的樣本，可以採用延後重取樣等方式。為了維護正確的輸入輸出關係（Input/X – Label/Y）和訓練是反向傳播的求導關係，實現中需要維護原始batch到每專家的sub-batch的index關係，在後來求導和結合合併時使用。
3. 獨立計算-Expert：併發（邏輯上可以先後）呼叫各個專家處理對應的sub-batch。這也是智慧平臺要支援的併發API之一。
4. 結果合併-Combine：合併每專家的結果tensor到整個batch的tensor，並按照資料分派索引，交換到原始輸入的順序。
在主流的深度學習智慧平臺中，可以採用兩類主要的實現策略：
張量置零：對需要分派到不同的後續網路單元（專家網路子網等），對需要分派的專家拷貝若干份tensor，對於不應輸入當前專家處理的資料維度置零。該方式在保證置零計算邏輯正確的情況下，實現簡單，全張量操作，對平臺無特殊要求，適用於演算法研究，僅體現條件計算前序資料被動態路由到不同的後續網路單元，分析演算法的效果。如果通過置零方式，該方法每個專家處理的tensor在batch維度大小是全batch，不能節省計算量和記憶體使用量。
張量整理：對需要分派到不同的後續網路單元（專家網路子網等），對需要分派的專家拷貝若干份tensor，對於不應輸入當前專家處理的資料維度不保留。並維護好sample級的index在變換前後的對應關係。在分散式友好的實現中，如果專家子網為單位被劃分到不同的計算節點，那麼專家網路的實現最好從子網級的平臺物件繼承後實現，比如：MindSpore中的mindspore.nn.Cell。詳細實現細節參見後續技術實現章節。
核心程式碼
核心程式碼：路由計算、資料分派、獨立計算，結果合併
參考程式碼採用MindSpore示意實現。(注：import mindspore as ms)
Mixture of Experts的核心邏輯，對輸入I，經過routing_network(最簡單*W即可)，然後topk(若變種演算法需要gate權重則需要softmax，否則可不)，然後用tensor的操作(可按照batch)選擇出每個subnetwork/expert的張量。
為方便除錯，採用了規模極小的非隨機的確定數值構造輸入和路由權重，路由網路採用簡單的X*W。
1、路由計算
當上述輸入5行（僅3類，希望分派給3個專家）樣本，和Gate權重做矩陣乘後，可以明確算出每個樣本要分派的專家。可以用matmul，也可以類似gates_weighted = einsum('bd,de->be', [data_inputs, gate_weights])第一輪矩陣乘的結果為：
輸入和權重乘法，在python中可以採用@，也可以採用matmul，也可以採用愛因斯坦求和簡記憶法函式einsum。當是簡單的矩陣乘的時候，採用einsum在計算圖編譯的時候實際會拆分成多個演算法，效能並不好；但當輸入和權重超過2維，需要以batch維固定做路由計算的時候，使用einsum可以程式設計實現很簡單。
2、分派
條件計算的分派，主要邏輯是根據路由網路的輸出，為每個樣本計算出top-k的專家。其實現可以通過topk函式實現。由於top選擇score可作為後續網路單元的輸入資訊（含路由的資訊），所以一般要對路由輸出做softmax做歸一化。
按需計算1：all-N專家之間的歸一化權重 (please refer to #2) ，gates_weighted一樣，按照dim=-1做了歸一化而已其輸出為：
為batch中每個sample選擇Top-K個專家 這裡為batch中每個的專家權重，可以從softmax-ed來top-k，也可以直接從gates_weighted來top-k；由於這裡可能不做softmax或者延後，所以可gates_weighted，這裡為batch中每個的專家序號
其輸出為：
接著：
按需計算2: top-n專家之間的歸一化權重
如何根據分派索引，從原始的輸入中，為每個專家提取出屬於該專家處理的tensor，在當前的主流智慧平臺，都沒有專門的運算元，可以通過其他運算元的組合來實現類似的效果。在MindSpore中，可以通過底層的C++實現運算元，也可以通過Python中繼承Cell並實現bprob， 然後將原始 gate tensor中按照index組織到目標輸出中。這裡我們實現一個Dispatch類
3、獨立計算
直接並行呼叫後續的專家網路。並行部分可以通過平臺來支援。可以通過特殊的函式或者annotation等標識，也可以由平臺編譯時優化為並行執行。（在非動態路由條件計算的網路模型中，一般不存在類似的優化。）
4、合併
合併的邏輯相對簡單，先通過cat按照batch維度做拼接，然後構造正確的zeros tensor用index_add按照索引將各個專家網路的結果在保持input序合併到一起，做為該MoE模組的輸出。
上述完成了整個MoE的完整計算過程。
程式碼框架
我們按照上述基本動態路由條件計算的張量操作為主的邏輯，擴充套件到一個完整的訓練程式碼框架中：
- class Dispatch(ms.nn.Cell): 實現路由中的分派邏輯
- class Combine(ms.nn.Cell): 實現路由中的組裝邏輯
- class Route(ms.nn.Cell): 完成整個動態路由邏輯，可以實現為相對通用的類
- class Expert(ms.nn.Cell): 平臺使用者自定義的專家網路
- class Network(ms.nn.Cell): 平臺使用者自定義的大網路
- class MSELoss(ms.nn.Cell)：實現MSE損失，實現輔助損失的邏輯
- class OutputLossGraph(ms.nn.Cell)：輸出infer和loss，PyNative模式單步
- class Dataset: 資料集類，僅滿足輸入shape和X-Y合理對應關係，僅僅示例def train( …): 訓練入口
條件計算實現技術點
1、動態路由
- 不可學習路由
如使用LSH (locality sensitive hashing)做路由：在整個可學習網路的前端，使用LSH來分派樣本，這樣可以避免LSH部分求導問題；如果在網路中間增加LSH模組，需要通過梯度估計完成確定性演算法部分梯度傳遞。
- 可學習路由
簡單的做法，定義gate_weights為可學習Parameter，對於二維的張量，通過python@或者matmul等完成權重路由計算；如果是更高維度的張量，且需固定batch維，einsum('bd*,*de->b*e')的形式完成計算。
2、topk和softmax的前後關係
在G_1(x)=softmax(topk(X*W)))和G_2(x)=topk(softmax(X*W)))兩類Gate實現中，
將softmax置於Topk前後，對top-k的選擇不變；當需要將G_*作為後序網路輸入的一部分，即將路由權重資訊作為後續網路輸入資訊，則需要考慮：需要all-N專家之間的歸一化權重，則softmax置於top-k之前；否則softmax置於top-k之後，來計算top-N專家之間的歸一化權重。
3、如何每專家在批次處理中平衡
按照每樣本的路由權重求和，即對batch單個樣本被分配的1+個export的重要性和權重求和，計算出importance；按照每樣本的路由權重中非0的求和，計算出有負載的專家來求得load。將coefficient_of_variation(importance) + coefficient_of_variation(load)作為auxiliary_loss參與優化，來平衡importance和load。變異係數(Coefficient of Variation)是用於無量綱度量資料的離散程度，越離散在此處表示均衡性越差，需要向更小優化。
在Transformer等多層（多處）MoE的模型中，將多組auxiliary_loss聯合作為auxiliary_loss, 在加dominated_loss之後即可。