選單

大規模深度學習模型預訓練技術一覽

深度學習是一種基於統計訓練的人工智慧演算法。深度神經網路模型往往可以透過預訓練方式來得到更高的效能。下面是讓我們來看看大型神經網路模型預訓練中的重要技術。

分散式並行

單一裝置的算力及記憶體容量受限於物理定律,已經難以滿足大模型規模擴大的需要,分散式訓練勢在必行。依據網路在分散式叢集中的切分方式,主要的分散式訓練模式包括資料並行、模型並行、流水並行,以及同時使用資料並行和模型並行的混合並行方法。誠然,分散式訓練通常會比單機單卡更快,且能支援更大規模的模型訓練。然而分散式訓練的不同模式也各有所長。在記憶體牆及網路牆的影響下,資料傳輸量的多寡成為影響分散式訓練的速度以及收斂性的關鍵因素。

資料並行

資料並行指叢集中的每個計算裝置上儲存相同的模型,但每個計算裝置處理不同的資料。訓練時,通常把每個批次的資料按份數均分到各計算裝置上進行處理。由於資料並行時每個裝置上都有完整的模型,所以每份資料在整個前向和後向計算過程中的中間結果都不需要資料傳輸,只是在更新模型時每一步都需要進行一次對模型引數的AllReduce操作。資料並行在引數量較小的模型上(如ResNet50)上表現優異,可以達到線性加速比。不同框架對資料並行的支援也趨於成熟,效能差異不大。

模型並行

模型並行則是把某一個運算元的模型引數劃分到多個裝置上去,每個裝置上只有這個運算元的一部分引數,所以每個批次的資料需要被傳輸到所有裝置上去處理。如果對基於運算元的輸出神經元進行引數切分,模型並行的每一步都需要輸入全部的資料,但是前向時的每一步輸出都只能得到部分資料。因此,為了使上一層網路的輸出能夠滿足下一層網路的輸入,每一步都需要對資料進行一次AllGather操作。在後向時對模型引數的更新不需要額外的傳輸操作,對資料梯度則需要ReduceScatter操作。資料並行與模型並行相比各有優劣。簡單來說,為了減少傳輸量,如果訓練時每個批次資料的中間結果數量較大則選擇資料並行;反之,若每個批次資料的計算中模型引數的數量較大則選擇模型並行。

流水並行

資料並行與模型並行都是讓裝置執行同一個層次的計算,而流水並行則是把任務劃分為幾個有明確先後順序的階段,把不同的階段分給不同的計算裝置,使得單裝置只負責網路中部分層的計算。在這種多裝置接力完成一個網路計算的模式下,可以支援更大的模型或者更大的批尺寸(batchsize)。由於是接力完成前向和後向計算,多裝置間只需要做資料的複製操作。另外,為了形成流水線,單裝置在完成一個批次的計算後立即開始下一個批次的計算,就需要將一個mini-batch拆分為多個micro-batch,或者模型更新部分使用非同步隨機梯度下降(SGD)機制。

混合並行

同一個神經網路的不同運算元可能適合不同的並行模式,但某個特定的運算元只能使用一種並行模式,例如在模型引數量大的地方使用模型切割,在模型引數量少的地方使用資料切割。相比於一個運算元只使用單一的並行模式,一個運算元同時使用多樣的並行模式可能進一步地減少傳輸量,譬如在隱藏層比較大的地方,同時對資料矩陣和模型矩陣進行切割。

大規模深度學習模型預訓練技術一覽

ZeRO-DP記憶體最佳化

資料並行與模型並行在記憶體及通訊量上各有侷限,前者在每個裝置上都會複製一份模型,後者需要在連續的計算之間插入資料傳輸,增大了掩蓋通訊開銷的難度。為此,微軟提出了ZeRO-DP技術,旨在克服兩種並行模式各自侷限的同時結合兩者的優點。其核心思想是,消除資料並行中每個裝置上都複製一份模型狀態(引數、梯度和最佳化器的狀態)所帶來的記憶體冗餘;把模型的狀態均攤到每個裝置上,每個裝置只需負責更新一部分引數,裝置程序間透過集合通訊操作同步梯度和引數,但是以單引數為粒度即用即獲取、用完即釋放。ZeRO-DP使得記憶體峰值佔用減少,單卡視訊記憶體大小的限制從整個模型狀態大小限制降低為一個分片模型狀態加上部分完整引數和梯度的大小限制。在一定程度上,該技術使得可訓模型的大小能隨著裝置數的增加而增加,且通訊開銷相比普通資料並行只增加了50%。不過ZeRO-DP仍假設了單計算裝置的記憶體能承載單個運算元的完整引數。當單個運算元的引數量超過了單裝置記憶體的容量時,ZeRO-DP就失效了,因此ZeRO-DP並不能完全取代模型並行。