如何基于MindSpore實(shí)現(xiàn)萬(wàn)億級(jí)參數(shù)模型算法？

本文是Switch Transformer的動(dòng)態(tài)路由條件計(jì)算的模型分析的第二篇 -?算法實(shí)現(xiàn)。

動(dòng)態(tài)路由條件計(jì)算的原理介紹可以參見上一篇《一文帶你了解MindSpore支持的萬(wàn)億級(jí)參數(shù)超大模型關(guān)鍵技術(shù)！》

實(shí)現(xiàn)策略

實(shí)現(xiàn)各種模型的帶有動(dòng)態(tài)路由稀疏激活的超大規(guī)模參數(shù)版本，需要分模型研究和實(shí)現(xiàn)。

(圖片來(lái)源：Switch Transformer論文)

3. 獨(dú)立計(jì)算-Expert：并發(fā)（邏輯上可以先后）調(diào)用各個(gè)專家處理對(duì)應(yīng)的sub-batch。這也是智能平臺(tái)要支持的并發(fā)API之一。

4. 結(jié)果合并-Combine：合并每專家的結(jié)果tensor到整個(gè)batch的tensor，并按照數(shù)據(jù)分派索引，交換到原始輸入的順序。

張量置零：對(duì)需要分派到不同的后續(xù)網(wǎng)絡(luò)單元（專家網(wǎng)絡(luò)子網(wǎng)等），對(duì)需要分派的專家拷貝若干份tensor，對(duì)于不應(yīng)輸入當(dāng)前專家處理的數(shù)據(jù)維度置零。該方式在保證置零計(jì)算邏輯正確的情況下，實(shí)現(xiàn)簡(jiǎn)單，全張量操作，對(duì)平臺(tái)無(wú)特殊要求，適用于算法研究，僅體現(xiàn)條件計(jì)算前序數(shù)據(jù)被動(dòng)態(tài)路由到不同的后續(xù)網(wǎng)絡(luò)單元，分析算法的效果。如果通過置零方式，該方法每個(gè)專家處理的tensor在batch維度大小是全batch，不能節(jié)省計(jì)算量和內(nèi)存使用量。

張量整理：對(duì)需要分派到不同的后續(xù)網(wǎng)絡(luò)單元（專家網(wǎng)絡(luò)子網(wǎng)等），對(duì)需要分派的專家拷貝若干份tensor，對(duì)于不應(yīng)輸入當(dāng)前專家處理的數(shù)據(jù)維度不保留。并維護(hù)好sample級(jí)的index在變換前后的對(duì)應(yīng)關(guān)系。在分布式友好的實(shí)現(xiàn)中，如果專家子網(wǎng)為單位被劃分到不同的計(jì)算節(jié)點(diǎn)，那么專家網(wǎng)絡(luò)的實(shí)現(xiàn)最好從子網(wǎng)級(jí)的平臺(tái)對(duì)象繼承后實(shí)現(xiàn)，比如：MindSpore中的mindspore.nn.Cell。詳細(xì)實(shí)現(xiàn)細(xì)節(jié)參見后續(xù)技術(shù)實(shí)現(xiàn)章節(jié)。

核心代碼

核心代碼：路由計(jì)算、數(shù)據(jù)分派、獨(dú)立計(jì)算，結(jié)果合并

Mixture of Experts的核心邏輯，對(duì)輸入I，經(jīng)過routing_network(最簡(jiǎn)單*W即可)，然后topk(若變種算法需要gate權(quán)重則需要softmax，否則可不)，然后用tensor的操作(可按照batch)選擇出每個(gè)subnetwork/expert的張量。

data_inputs = ms.Tensor([ [0.1,0.9], [0.8,0.8], [0.9,0.1], [0.1,0.9], [0.9,0.1], ]) #假設(shè)輸入為5個(gè)樣本，每個(gè)2維，當(dāng)然可以擴(kuò)展到高維 (batch,dimension) = (5,2) gate_weights = ms.Parameter(ms.Tensor([ [0.1,0.5,0.9], [0.9,0.5,0.1], ] , ms.float32) , requires_grad=True) #假設(shè)路由門權(quán)重，3個(gè)專家，每個(gè)2維和輸入一樣 (dimension,experts) = (2,3)

gates_weighted= [[0.8200, 0.5000, 0.1800], [0.8000, 0.8000, 0.8000], [0.1800, 0.5000, 0.8200], [0.8200, 0.5000, 0.1800], [0.1800, 0.5000, 0.8200]]

gates_softmax = softmax(input=gates_weighted, axis=-1)

gates_softmax= [[0.4438, 0.3222, 0.2340], [0.3333, 0.3333, 0.3333], [0.2340, 0.3222, 0.4438], [0.4438, 0.3222, 0.2340], [0.2340, 0.3222, 0.4438]]

為batch中每個(gè)sample選擇Top-K個(gè)專家 這里為batch中每個(gè)的專家權(quán)重，可以從softmax-ed來(lái)top-k，也可以直接從gates_weighted來(lái)top-k；由于這里可能不做softmax或者延后，所以可gates_weighted，這里為batch中每個(gè)的專家序號(hào)

gates_topk_value, gates_topk_index =topk(gates_softmax, 1)

其輸出為：

gates_softmax= [[0.4438, 0.3222, 0.2340], [0.3333, 0.3333, 0.3333], [0.2340, 0.3222, 0.4438], [0.4438, 0.3222, 0.2340], [0.2340, 0.3222, 0.4438]]

接著：

gates_topk_value, gates_topk_index =topk(gates_softmax, 1)

按需計(jì)算2: top-n專家之間的歸一化權(quán)重

class Dispatch(ms.nn.Cell): def __init__(self, expert_number): super().__init__() self.expert_number = expert_number self.reshape = ms.ops.Reshape() self.concat = ms.ops.Concat() self.zeros = ms.ops.Zeros() self.add = ms.ops.AddN() def set_indices_in(self, indices_in): #可以作為construct的參數(shù) self.indices_in = indices_in def get_indices_out(self): #可以用construct的返回值返回 return self.indices_out def construct(self, data): dispatch = [] indices_out = [] for _ in range(self.expert_number): dispatch.append([]) indices_out.append([]) for uid,(idx,dat) in enumerate(zip(self.indices_in, data)): dat = self.reshape(dat, (1, dat.shape[0])) if len(dispatch[idx]) == 0: dispatch[idx] = dat indices_out[idx] = [uid] else: dispatch[idx] = self.concat((dispatch[idx], dat)) indices_out[idx] = indices_out[idx]+[uid] self.indices_out = [y for x in indices_out for y in x] return dispatch def bprop(self, data, out, dout): #反向梯度計(jì)算 dall = None for one in dout: if dall == None: dall = one else: dall = self.concat((dall, one)) do = self.zeros(dall.shape, ms.float32) for idx_target, idx_source in enumerate(self.indices_out): do[idx_target] = self.add((do[idx_target], dall[int(idx_source)])) return do

直接并行調(diào)用后續(xù)的專家網(wǎng)絡(luò)。并行部分可以通過平臺(tái)來(lái)支持。可以通過特殊的函數(shù)或者annotation等標(biāo)識(shí)，也可以由平臺(tái)編譯時(shí)優(yōu)化為并行執(zhí)行。（在非動(dòng)態(tài)路由條件計(jì)算的網(wǎng)絡(luò)模型中，一般不存在類似的優(yōu)化。）

4、合并

合并的邏輯相對(duì)簡(jiǎn)單，先通過cat按照batch維度做拼接，然后構(gòu)造正確的zeros tensor用index_add按照索引將各個(gè)專家網(wǎng)絡(luò)的結(jié)果在保持input序合并到一起，做為該MoE模塊的輸出。

class Combine(ms.nn.Cell): def __init__(self): super().__init__() self.zeros = ms.ops.Zeros() self.add = ms.ops.AddN() def set_indices(self, indices): #可以作為construct的參數(shù) self.indices = indices def construct(self, data): O = self.zeros(data.shape, ms.float32) for idx_target, idx_source in enumerate(self.indices): O[idx_target] = self.add((O[idx_target], data[int(idx_source)])) return O def bprop(self, data, out, dout): #反向梯度計(jì)算 do = self.zeros(dout.shape, ms.float32) for idx_target, idx_source in enumerate(self.indices): do[idx_target] = self.add((do[idx_target], dout[int(idx_source)])) return do

上述完成了整個(gè)MoE的完整計(jì)算過程。

代碼框架

條件計(jì)算實(shí)現(xiàn)技術(shù)點(diǎn)

1、動(dòng)態(tài)路由

不可學(xué)習(xí)路由

如使用LSH (locality sensitive hashing)做路由：在整個(gè)可學(xué)習(xí)網(wǎng)絡(luò)的前端，使用LSH來(lái)分派樣本，這樣可以避免LSH部分求導(dǎo)問題；如果在網(wǎng)絡(luò)中間增加LSH模塊，需要通過梯度估計(jì)完成確定性算法部分梯度傳遞。

可學(xué)習(xí)路由

簡(jiǎn)單的做法，定義gate_weights為可學(xué)習(xí)Parameter，對(duì)于二維的張量，通過python@或者matmul等完成權(quán)重路由計(jì)算；如果是更高維度的張量，且需固定batch維，einsum('bd*,*de->b*e')的形式完成計(jì)算。

2、topk和softmax的前后關(guān)系

在G_1(x)=softmax(topk(X*W)))和G_2(x)=topk(softmax(X*W)))兩類Gate實(shí)現(xiàn)中，

將softmax置于Topk前后，對(duì)top-k的選擇不變；當(dāng)需要將G_*作為后序網(wǎng)絡(luò)輸入的一部分，即將路由權(quán)重信息作為后續(xù)網(wǎng)絡(luò)輸入信息，則需要考慮：需要all-N專家之間的歸一化權(quán)重，則softmax置于top-k之前；否則softmax置于top-k之后，來(lái)計(jì)算top-N專家之間的歸一化權(quán)重。

3、如何每專家在批次處理中平衡

按照每樣本的路由權(quán)重求和，即對(duì)batch單個(gè)樣本被分配的1+個(gè)export的重要性和權(quán)重求和，計(jì)算出importance；按照每樣本的路由權(quán)重中非0的求和，計(jì)算出有負(fù)載的專家來(lái)求得load。將coefficient_of_variation(importance) + coefficient_of_variation(load)作為auxiliary_loss參與優(yōu)化，來(lái)平衡importance和load。變異系數(shù)(Coefficient of Variation)是用于無(wú)量綱度量數(shù)據(jù)的離散程度，越離散在此處表示均衡性越差，需要向更小優(yōu)化。

了解完MindSpore的關(guān)鍵技術(shù)是不是很心動(dòng)呢！趕緊【點(diǎn)擊鏈接】并【立即報(bào)名】，即可在 ModelArts 平臺(tái)學(xué)習(xí)到一個(gè)經(jīng)典案例掌握基于MindSpore的深度學(xué)習(xí)！

想要了解更多關(guān)于大模型的知識(shí)，請(qǐng)點(diǎn)擊：專家解惑 | 關(guān)于華為云盤古大模型，你想問的都在這里~

實(shí)實(shí)現(xiàn)實(shí)現(xiàn)策略策略現(xiàn)策略

MindSpore 專家

版權(quán)聲明：本文內(nèi)容由網(wǎng)絡(luò)用戶投稿，版權(quán)歸原作者所有，本站不擁有其著作權(quán)，亦不承擔(dān)相應(yīng)法律責(zé)任。如果您發(fā)現(xiàn)本站中有涉嫌抄襲或描述失實(shí)的內(nèi)容，請(qǐng)聯(lián)系我們jiasou666@gmail.com 處理，核實(shí)后本網(wǎng)站將在24小時(shí)內(nèi)刪除侵權(quán)內(nèi)容。

版權(quán)聲明：本文內(nèi)容由網(wǎng)絡(luò)用戶投稿，版權(quán)歸原作者所有，本站不擁有其著作權(quán)，亦不承擔(dān)相應(yīng)法律責(zé)任。如果您發(fā)現(xiàn)本站中有涉嫌抄襲或描述失實(shí)的內(nèi)容，請(qǐng)聯(lián)系我們jiasou666@gmail.com 處理，核實(shí)后本網(wǎng)站將在24小時(shí)內(nèi)刪除侵權(quán)內(nèi)容。