面向混合異構架構的模型并行訓練優化方法*

高 開，郭振華，陳永芳，王 麗，趙雅倩，趙 坤

(1.浪潮電子信息產業股份有限公司高效能服務器與存儲技術國家重點實驗室，山東 濟南 250000； 2.廣東浪潮大數據研究有限公司，廣東 廣州 510000)

1 引言

深度神經網絡(DNN)推進諸多機器學習任務的進步，其中包括語音識別、視覺識別和語言處理等。BigGan、Bert和GPT2.0 等最新成果表明，DNN 的模型越大，任務處理的表現就越好，而該結論也在過去的視覺識別任務中得到了驗證，這表明模型大小與分類準確性之間存在很強的關聯性。與此對應的是，各個公司推出了各種類型的加速器設備用來加速訓練神經網絡模型，Intel和XILINX分別推出了高效能的FPGA設備，Google推出了一系列的TPU(Tensor Processing Unit)加速板卡。目前最先進的圖像模型早已達到 TPUv2 可用內存的極限，單一類型的板卡已經無法用來訓練很大的模型，因此，迫切需要一種高效、可擴展的基礎設施，以實現大規模的深度學習訓練，并克服當前加速器內存受限問題，進而可以擴展到多種類型的加速器設備上，能夠充分使用混合異構集群中存在的多種類型且性能不一的設備。

為了在多個計算設備之間部署深度學習DL(Deep Learning)模型以訓練大型且復雜的模型，數據并行DP(Data Parallel)是應用最廣泛的并行化策略，但是如果數據并行訓練中的設備數量不斷增加，設備的性能不一，設備之間的通信開銷就會成為瓶頸。此外，對于性能分布不均的設備群，訓練時每個步驟都需要給性能不同的設備分配不同的批處理大小，這種不均衡會導致精度損失，即需要更長的訓練周期以收斂到期望的精度[1]。這些因素會影響數據并行整體訓練時間，并且會影響某些設備的運行效率。

除了DP，還可以利用模型并行性MP(Model Parallel)來加速模型訓練。模型并行首先把整個模型按照一定的策略劃分到各個設備上，然后通過設備間的相互通信來解決特征數據傳輸問題。Gpipe[2]是一種典型的模型并行訓練框架，它根據設備的性能和模型每一層的計算量大小按照相應比例為各個計算設備配置模型，然后依次按照設備順序并行流水地進行訓練。這種方法可解決數據并行中的精度損失問題，但同樣會因為設備群中設備數量增加帶來流水線過長和通信瓶頸問題。同時，在混合異構平臺中，模型并行訓練時需要更加精細地劃分模型，因此如何在性能分布不均的設備上劃分模型，以及如何高效地建立設備間的流水都成為了重要的研究熱點問題[3]。

2 基于Gpipe框架的模型并行訓練

簡單來說，Gpipe 是一款分布式機器學習庫，基于同步隨機梯度下降與流水并行技術進行模型并行訓練，適用于由多個連續層組成的任意 DNN。最重要的是，Gpipe 可以使研究人員在不調整超參數的情況下部署更多的加速器，以訓練大規模模型，由此有效擴展了性能。

為了實現跨加速器的高效訓練，Gpipe 先按照加速器對模型進行按層劃分，然后將每個批次的訓練示例拆分為更小的微批次。通過在微批次中執行流水管理，加速器得以并行流水運行。此外，梯度將在微批次中持續累積，以免分區的數量影響到模型的質量。

Gpipe框架中采用的模型并行策略是按照模型每層的計算量均等地分到各個計算設備上(如圖1a所示)，然后在各個設備上依次進行模型的前向反向計算(如圖1b所示)。從圖1b可以看到，每個時刻只有一個設備在執行計算任務，其他設備都處在空閑狀態，這樣會造成極大的資源浪費。為了解決這個問題，Gpipe通過把每一個批次的數據劃分成更小批次的方法，使得每個小批次在設備間流水執行，盡可能地使所有設備都能夠同時并行計算(如圖1c所示)。Gpipe一般是通過總的批處理總量除以設備數來得到更小的批次。對于具有單一類型設備的異構平臺，Gpipe的模型并行劃分策略和微批次流水的方法都具有較好的性能。但是，對于混合異構平臺，由于一些設備性能分布差距過大，均等的模型劃分策略和微批次劃分方法都不能充分利用集群中設備的性能。

Figure 1 An example of Gpipe working principle圖1 Gpipe 工作原理示例

圖1a是模型并行訓練中序列模型劃分示例。其中，Fk為第k個部分的前向計算，Bk為第k個部分的后向計算，依賴于Bk+1的輸出和Fk的輸出。圖1b是傳統的模型并行訓練的前向反向過程。由于網絡存在順序性，每次都只有一個加速器處于活動狀態。 模型并行策略導致嚴重的設備利用不充分問題。
圖1c是Gpipe 將輸入的小批次拆成更小的微批次，使不同的加速器可以同時在單獨的微批次上運行。

3 混合異構架構下模型并行劃分映射策略

Gpipe所提供的模型并行流水訓練一般適用于性能均等的加速設備，比如Google在一系列TPUv2芯片上實現了模型的劃分和并行流水訓練[4]。Gpipe中使用的模型劃分策略是按照模型中每一層計算量和設備的計算性能進行均等劃分。而對于混合異構平臺，通常每臺設備上都部署了多種類型的加速芯片，比如GPU、FPGA和TPU設備等。當設備群中出現性能不一的多種類型設備時，每種設備的計算能力和通信帶寬的不一致給網絡模型的劃分帶來很大的問題[5]，網絡模型就不能簡單地按層進行劃分。模型并行訓練時需要同時利用多種類型的加速芯片，如何在多種性能分布不均的設備上部署模型以及合理地建立設備間的流水是急需解決的問題。

在混合異構架構下，設備的性能通常會有很大的差異性，比如設備群的FPGA設備(如圖2a所示)，通常情況下FPGA設備只能承擔某一層的計算量，甚至某些性能過低的FPGA設備無法承擔一層的計算量，因此需要更加精細的模型并行訓練時模型劃分方法。為了解決設備性能分布不均帶來的模型劃分困難，本文提出了一種層級并行和通道并行混合的模型劃分策略。

Figure 2 Model parallel training method based on Gpipe圖2 基于Gpipe框架的層級模型并行

本文以混合異構平臺中有2個GPU設備和2個FPGA設備為例說明層級并行和通道并行的模型混合并行策略。由于使用的FPGA設備的性能要低于GPU設備的性能，首先把2個FPGA設備合并成1個設備，然后在合并后的3個設備上用傳統的層級劃分策略進行模型劃分，即按照每層的計算量和設備的計算力按比例把網絡模型按層劃分到每個設備上(如圖3a所示)。按層級劃分完成模型劃分之后，再把劃分到2個FPGA設備上的網絡層按卷積核通道數在2個FPGA設備上按計算能力進行均等劃分(如圖3b所示)。如果2個FPGA設備的性能均等，那么劃分到每個設備上的通道數是一致的。

Figure 3 Model parallel training of hierarchy parallel and channel parallel圖3 層級并行與通道并行混合的模型并行

這種混合層級并行和通道并行的劃分策略有2個方面的優點：(1)混合并行策略對網絡模型的劃分是更加精細的，模型的劃分粒度不僅僅局限在以層級為單位的劃分，而是精細到了以通道為單位，這種混合劃分策略能夠更充分地利用不同類型設備的性能。(2)混合并行劃分策略可以減少混合異構中設備流水線的長度(如圖2c和圖3c所示)，為第4節模型并行訓練時流水線的建立提供了便利。

Figure 4 Idle time of model parallel pipeline training when device number is 3 and micro-batch-size is 4圖4 模型并行訓練流水過程中設備空閑時間(設備數量為3，微批次數量為4)

4 流水線的優化與微批次的自動調節

Gpipe框架是通過把批次劃分成更小的微批次來實現設備間的并行流水訓練，劃分后的批次大小B=M/K，其中M為原批次大小，K為設備數量。圖4闡明了設備數量K為3，微批次數量為4時設備間流水效果。

通過圖4的流水過程分析可以得到，訓練過程中有一半以上的時間至少有1個設備處于空閑狀態，因此這樣的流水效果沒有充分利用設備的性能，反而在流水線建立上的時間占比很大[6]。為了縮短流水線建立的時間，傳統的模型并行算法通常是通過劃分更小的微批次來掩蓋流水線建立的時間。一般情況下，當劃分的微批次數量越多，設備間的流水效果會更好，當劃分后的微批次數量N>4×K，K為設備數量，流水線的建立時間就會占比很小。同時，如果混合異構平臺中設備數量巨大，就需要劃分更多的微批次來達到較好的流水效果。顯而易見的是，批處理總量是不變的，當劃分的微批次數量越多，每次每個設備處理的批量大小就會越小。若設備每次處理的批量大小遠小于設備的處理能力的時候，通常也會無法充分利用該設備的性能資源。

如何根據流水線長度和設備的處理能力來確定微批次的大小和數量就是一個急需解決的問題，本文根據以下2個原則來確定如何選擇微批次的劃分數量。

4.1 原則1：化零為整，減小流水線的長度

在模型并行訓練中，當設備數量增多時，通常流水線的長度也會隨之增加，因此花費在流水建立的時間占比也會增大。而且混合異構架構下設備的類型多種多樣，通常會有很大的性能差異性(如圖5所示)，這些差異性會給設備間流水的建立帶來通信瓶頸問題。

為了解決上述問題，本文提出一種混合異構架構下多設備合并的策略，通過合并一些性能較低的設備使合并后的設備的性能分布均勻(如圖6所示)。設備合并采用通過循環遍歷設備性能然后合并性能較低的2個設備的方式。假設第i個設備的性能記為Pi，整個設備群中設備的性能分布用平均差來刻畫，如式(1)所示：

(1)

算法1設備合并方法device_merge()

輸入：設備性能集合{Pi,…,Pn}。

輸出：合并后的設備性能集合{Pi,…,Pk}。

while(1)

{

{Pi,…,Pn}=Sort({Pi,…,Pn});

Pn-1=Pn-1+Pn;

Output {Pi,…,Pn-1};

Break;

}

n=n-1;

}

經過設備合并之后的設備群流水線長度縮短了(如圖6所示)，各設備性能差別很小或者已經沒有差別。在完成設備合并之后需要進行模型的劃分和映射，采用3.1節中的層級并行和通道并行混合劃分策略，首先根據合并之后設備的性能和模型各層的計算量按比例進行劃分；然后在合并的設備上根據通道數目和設備性能按比例進行劃分。本文提出的這種設備合并方法能夠和第3節提到的混合并行策略完美地結合在一起，不僅縮短了流水線的長度，而且還緩解了設備群中的性能較低設備的通信壓力。

Figure 5 The pipeline training before device combined when device number is 5 and micro-batch-size is 4圖5 設備合并前的流水效果圖(設備數量為5，微批次數量為4)

Figure 6 The pipeline training after device combined when device number is 5 and micro-batch-size is 4圖6 設備合并后的流水效果圖(設備數量為5，微批次數量為4)

4.2 原則2：負載均衡，自動調節微批次

在經過上述設備合并之后模型并行訓練需要建立設備間的流水，本文采用Gpipe框架中微批次劃分的方法使得各小批次在設備間流水執行。Gpipe框架中微批次劃分方法是總的處理批次大小除以設備數量。這種方法通常不能完全利用設備性能或者流水建立的時間占比很大，為了緩解這個問題，本文提出根據流水線的長度和設備性能來動態地進行微批次劃分。首先，計算得到合并后每個設備的傳輸帶寬和計算能力；然后分析批次的劃分對流水建立時間占比和設備利用率的影響。圖7是2個設備時微批次的劃分數量大小對流水建立時間占比和設備性能利用率的影響曲線。

Figure 7 Influence of pipelining bubble time and device effective with the number of micro-batch圖7 微批次的劃分數量對流水建立時間和設備利用率的影響

假設流水建立時間對整個訓練時間的影響為Ppipe(batch_num)，設備利用率損失為Dloss(batch_num)，則batch_num需要滿足以下條件：

min(Ppipe(batch_num)+Dloss(batch_num))

(2)

滿足式(2)的batch_num能夠使設備間的流水效果和設備的利用率達到一個均衡的狀態，進而整個系統的性能達到最優狀態。本文采用二分查找方法找到batch_num最優解的。假設總的處理批量大小為Batch_size，混合異構平臺中的設備最大處理批次為Batch_device，流水線長度為Pipe_num。流水線建立的時間占比為：

(3)

設備利用率損失為：

(4)

通常式(4)計算得到batch_num，劃分后的微批次由式(5)計算得到：

(5)

為了防止計算得到的micro_batch不為整數，對每個micro_batch進行±1微調，使其都為整數且相加等于總的批處理大小。

表1以ResNet50模型為例說明了流水線的長度與Batch_size如何進行選擇。每次處理的總批量為256，混合異構平臺中以8個GPU為例，每個GPU處理的最大Batch_size為64。通過二分法查找得到表1的劃分結果。

Table 1 Result of micro-batch on ResNet50 with hybrid heterogeneous devices表1 混合異構架構下ResNet50的劃分結果

5 實驗

首先，本文對比在混合異構平臺上模型并行算法與數據并行算法在精度上的差別，以標準的ResNet50模型在單個GPU上訓練的精度(圖8中的Original)為基準，分別對比了數據并行和模型并行在混合異構平臺中的表現。實驗采用2種不同型號(2個NVIDIA Tesla V100和2個NVIDIA Tesla M4)的GPU設備，這2種型號的GPU設備在性能上有一定的差別。實驗環境如下所示：

CPU：2路Intel(R)Xeon(R) CPU E5-2690 v3 @ 2.60 GHz(每個CPU包含12個物理core，24個thread cores);FPGA：Inspur F10S板卡x1;GPU:NVIDIA Tesla V100 ×2 NVIDIA Tesla M4 ×2;RAM：629 GB;OS：CentOS Linux release 7.6.1810 (Core);Kernel version:3.10.0-514.el7.x86_64;輸入圖像：ImageNet2012。

為了充分利用每個設備的性能，在數據并行時，給不同型號的GPU設備設置的Batch_size是不一樣的，例如，NVIDIA Tesla V100每次處理的Batch_size設置為128，NVIDIA Tesla M4每次處理的Batch_size設置為64。在使用本文模型進行并行訓練時，首先把2個NVIDIA Tesla M4合并成一個設備，使流水線長度減小為3；然后用層級并行和通道并行的混合并行策略把ResNet50模型分配到各個設備上;最后選取合適的微批次劃分(32×4)，使整個設備群的流水狀態達到平衡。通過實驗發現，模型并行訓練方法可以彌補混合異構架構中數據并行訓練時出現的精度損失問題，如圖8所示，數據并行訓練算法訓練過程前期精度波動較大，通常需要更多的訓練周期才能達到相同的精度Top-1 accuracy，如圖9所示。

Figure 8 Accuracy of data parallel and model parallel on hybrid heterogeneous platform圖8 混合異構架構下數據并行算法 與模型并行算法的精度(Epoch=90)

Figure 9 Accuracy curve of data parallel and model parallel on hybrid heterogeneous platform圖9 混合異構下精度收斂曲線對比

然后，與Google提出的Gpipe模型并行流水算法在單異構和混合異構平臺上的性能表現進行對比。與傳統的模型并行流水算法(Gpipe)相比，通過本文方法優化之后的模型并行流水算法性能得到了相應的提升。本文基于Gpipe與TensorFlow實現了ResNet50的模型并行流水訓練，并且成功地部署到CPU、GPU和FPGA設備上進行訓練。本文分別在完全由GPU組成的單異構訓練環境和由GPU與FPGA共同組成的混合異構訓練環境進行了實驗驗證。通過本文模型劃分策略和批次大小調整以后，單異構環境下ResNet50模型的訓練性能比Gpipe框架提升了3%～4%,如圖10所示。對于由GPU和FPGA組成的混合環境，由于FPGA設備性能的限制，只把模型的最后2層放到了FPGA設備上。實驗表明，混合異構環境下經過本文提出的混合并行策略和批次劃分方法優化之后，比傳統的模型并行流水算法的性能加速比提升了7%左右，如圖11所示。

Figure 10 Speedup of Gpipe and our method on single heterogeneous platform圖10 單異構環境下性能加速比

Figure 11 Speedup of original and our method on hybrid heterogeneous platform圖11 混合異構環境下性能加速比

6 結束語

本文提出了2種適用于混合異構架構平臺的模型并行訓練優化方法。首先針對混合異構平臺中設備性能的差異性提出了混合層并行和通道并行的模型劃分策略，相比簡單的層并行方法，這種策略對模型的劃分更加精細，能夠更加充分地利用設備性能。然后本文給出了一種動態劃分微批次的方法，該方法中的2個原則能夠使設備間流水線的建立和設備利用率達到平衡，進而能夠最大限度地減小流水線建立的時間占比和提高設備的利用率。相比于數據并行方法，本文提出的模型并行方法能夠解決數據并行在混合異構架構中因設備性能分布不均帶來的精度損失問題。相比傳統的模型并行算法(Gpipe)具有更好的加速比，在單一類型設備的異構平臺上的訓練性能加速比提升4%左右，在混合異構平臺的訓練性能加速比相比沒有使用優化方法之前能夠提升7%左右。