 作者 馭策（龔志剛） 筍江（林立翔）蜚廉（王志明） 昀龍（游亮） 近些年來，深度學(xué)習(xí)在圖像識(shí)別，自然語言處理等領(lǐng)域快速發(fā)展。各種網(wǎng)絡(luò)模型，需要越來越多的計(jì)算力來進(jìn)行訓(xùn)練。以典型的中等規(guī)模的圖像分類網(wǎng)絡(luò)Resnet50為例，基準(zhǔn)的訓(xùn)練精度為Top-1 76%， Top-5 為 93%，為達(dá)到此精度，一般需要將整個(gè)Imagenet數(shù)據(jù)集的128萬張圖片，訓(xùn)練90次(90 epoch). 這樣的計(jì)算量，以單張P100的計(jì)算力需要6天才能訓(xùn)練完畢，而最近的NLP 領(lǐng)域取得突破的Bert 模型以及GPT-2的預(yù)訓(xùn)練，如果在單機(jī)上進(jìn)行則需要數(shù)月甚至按年計(jì)。因此在實(shí)際的生產(chǎn)環(huán)境，引入分布式訓(xùn)練，大大降低模型訓(xùn)練所需的時(shí)間，提高模型迭代的速度成為緊迫的需求。基于這樣的需求，各大深度學(xué)習(xí)框架，基本都實(shí)現(xiàn)了分布式訓(xùn)練的支持。  在實(shí)踐中，我們發(fā)現(xiàn)在云環(huán)境中，現(xiàn)有訓(xùn)練框架的分布式有以下問題： 不同的框架支持分布式的具體形式均不相同，支持用戶從單機(jī)到分布式的工作量和框架的種類成正比。 所有框架的分布式實(shí)現(xiàn)，在云上的擴(kuò)展性均不理想，初步的profiling結(jié)果顯示，網(wǎng)絡(luò)帶寬并未被有效利用。以Resnet50為例，Tensorflow + Horovod 四機(jī)32卡，batch size為128，GPU卡為1080Ti，擴(kuò)展性僅為86%。在VPC25Gb網(wǎng)絡(luò)下，這個(gè)擴(kuò)展性明顯低于預(yù)期。 為了進(jìn)行擴(kuò)展性的優(yōu)化，需要考慮對(duì)諸多不同框架的支持，面臨巨大的重復(fù)性工作且難于維護(hù)。 以上3個(gè)問題都指向了一個(gè)方向，就是統(tǒng)一各框架的分布式通信到一個(gè)單一的架構(gòu)中。我們也注意到開源的Horovod項(xiàng)目，部分解決了以上問題，目前已經(jīng)支持Tensorflow，Pytorch和Keras，最近也加入了MXNet的支持。Horovod給了我們一個(gè)很好的參考，這種大統(tǒng)一的分布式通信框架是可行的，梯度的同步通信是完全可以被抽象為框架無關(guān)的架構(gòu)的。但是Horovod最大的問題是沒有解決云上的擴(kuò)展性問題，我們經(jīng)過討論，決定開發(fā)一個(gè)完全獨(dú)立的項(xiàng)目Perseus，在阿里云上打造一個(gè)高度優(yōu)化的統(tǒng)一深度學(xué)習(xí)分布式通信框架。 通用分布式通信架構(gòu)設(shè)計(jì) 基于阿里云的基礎(chǔ)設(shè)施的特點(diǎn)，我們?cè)O(shè)計(jì)了一個(gè)完全自主的通用分布式框架，除了支持Horovod所支持的Tensorflow和Pytorch以外，我們根據(jù)阿里云的主要AI用戶的需求，還增加了MXNet以及Caffe的支持。系統(tǒng)的架構(gòu)如下圖：  Perseus 訓(xùn)練通信框架整體模塊圖 系統(tǒng)設(shè)計(jì)的宗旨是將所有與實(shí)際通信相關(guān)的操作全部對(duì)上層訓(xùn)練框架保持透明，同時(shí)將對(duì)不同框架的支持切分為互相獨(dú)立的模塊。基于這樣的宗旨，具體模塊設(shè)計(jì)如下： 接口層 通信框架需要提供兩類接口，一類為Python 接口，另一類為C接口。接口層的任務(wù)是對(duì)上層模型訓(xùn)練任務(wù)提供以下幾個(gè)功能： 創(chuàng)建通信層以及對(duì)應(yīng)的框架支持 獲取當(dāng)前進(jìn)程的全局任務(wù)編號(hào)，全局任務(wù)規(guī)模，節(jié)點(diǎn)內(nèi)任務(wù)編號(hào)，和節(jié)點(diǎn)內(nèi)任務(wù)規(guī)模 關(guān)閉當(dāng)前通信框架 為MXNet，Pytorch和Tensorflow需要準(zhǔn)備對(duì)應(yīng)功能的python接口，為Caffe只需要提供C的接口。各框架的Python接口部分的細(xì)節(jié)描述，見框架層。 通信層
每個(gè)框架在進(jìn)行梯度歸約通信之前，必須先注冊(cè)梯度，注冊(cè)過程的參數(shù)是梯度的上層訓(xùn)練框架層的唯一標(biāo)識(shí)的名稱，內(nèi)部會(huì)分配一個(gè)key值來代表改梯度；之后就可以進(jìn)行梯度通信，接口是將對(duì)應(yīng)的梯度加入到通信隊(duì)列之中，并標(biāo)明通信類別為廣播或者是Allreduce，同時(shí)還需要注冊(cè)一個(gè)回調(diào)函數(shù)，當(dāng)分布式通信結(jié)束后，會(huì)調(diào)用回調(diào)函數(shù)通知上層訓(xùn)練框架梯度通信已經(jīng)完成，可以進(jìn)行后續(xù)的參數(shù)更新操作了。梯度通信的所有細(xì)節(jié)，都不需要暴露任何接口。 通信層內(nèi)部模塊圖如下： Perseus 通信模塊圖 通信的主要工作都是在后臺(tái)主線程中完成，異步的等待在多流線程池中進(jìn)行，這樣的完全異步多流方式，將可能的單一慢連接對(duì)全局的影響降到最低。主要的步驟如下： 主線程中處理等待通信的梯度 多節(jié)點(diǎn)協(xié)商各梯度的狀態(tài)，梯度是否就緒 確定各節(jié)點(diǎn)產(chǎn)生順序完全一致的就緒梯度列表 各節(jié)點(diǎn)統(tǒng)計(jì)各自列表是否符合通信粒度需求 開始進(jìn)行融合分片通信，這里需要處理與上層框架的同步關(guān)系，確保開始融合通信之前，上層框架已經(jīng)產(chǎn)生對(duì)應(yīng)的梯度輸出。分片的多流資源是同步獲取，若無可用流資源，則等待之前的某些流結(jié)束。若有流資源，則將該分片數(shù)據(jù)，使用該流發(fā)出，同時(shí)通知線程池中對(duì)應(yīng)線程。 完成梯度的分散，對(duì)應(yīng)片斷的回調(diào)函數(shù)，在片斷回調(diào)函數(shù)中會(huì)檢查是否已經(jīng)完整整個(gè)梯度，若是則通知上層訓(xùn)練框架，該梯度已經(jīng)通信完成。 多流線程池 每個(gè)流有一個(gè)專用的線程，會(huì)等待主線程的通知。有通知后，等待NCCL 對(duì)應(yīng)的cuda 流結(jié)束，確保上半部通信已經(jīng)完成，生成下半部操作，加入對(duì)應(yīng)隊(duì)列。 在通信層中，對(duì)真正進(jìn)行通信的更底層的原語進(jìn)行了封裝，目前支持下面兩類通信模式：
框架層 Tensorflow 與Pytorch 對(duì)Tensorflow與Pytorch的支持，考慮到Horovod在社區(qū)的影響力，為了方便用戶已有的模型代碼能夠無縫的整合到Perseus 通信庫，經(jīng)過討論，對(duì)這兩種框架的框架層支持的python接口部分，選擇了與Horovod完全一致的接口，讓用戶的切換成本接近于零。而與框架層本身銜接的部分，需要定義各自的Tensor以及與Context的虛類的實(shí)現(xiàn)即可，均不復(fù)雜。 實(shí)現(xiàn)完畢之后，對(duì)于之前使用Horovod的分布式模型，只需要進(jìn)行下面的替換即可：  MXNet MXNet 框架的支持細(xì)節(jié)可以參考Perseus MXNet這篇ATA文章，這里不再重復(fù)描述。 Caffe 由于Caffe本身的分布式支持比較原始，而且沒有模塊化，因此Caffe的支持是這四類框架相對(duì)困難的部分。前面3類框架的支持，除了MXNet需要對(duì)MXNet做少量修正外，其余的均不需要改變框架的任何代碼，但是對(duì)Caffe，我們需要對(duì)Caffe的框架做不小的修改，這些修改主要有以下幾個(gè)方面 將單Process 多GPU卡的模式改造為單Process單GPU卡，以MPI的方式來launch 多機(jī)多卡的訓(xùn)練 將梯度融合部分修改為使用Perseus的框架API來完成 需要構(gòu)造callback機(jī)制，讓Perseus通信框架能夠通知Caffe框架整個(gè)batch的梯度均已通信完畢。 Perseus的Caffe框架部分也要增加相應(yīng)的實(shí)現(xiàn)。最后整合才得以實(shí)現(xiàn)Caffe的多機(jī)多卡。 通信的優(yōu)化 實(shí)現(xiàn)了以上架構(gòu)后，我們成功的把所有訓(xùn)練框架的分布式通信部分整合到了同一個(gè)框架之中。接下來在這套通信框架之下的所有優(yōu)化工作，都將直接使所有的訓(xùn)兩框架受益。我們第一階段的大部分優(yōu)化工作都和VPC網(wǎng)絡(luò)深度相關(guān)，這部分的優(yōu)化總體可以分為兩個(gè)類別 帶寬優(yōu)化 我們通過對(duì)Horovod的NCCL實(shí)現(xiàn)方式進(jìn)行了Profiling，發(fā)現(xiàn)很難達(dá)到云上網(wǎng)絡(luò)帶寬的上限，經(jīng)過分析我們確認(rèn)，主要的原因是TCP網(wǎng)絡(luò)，單流通信無法達(dá)到上限，使用Iperf也證實(shí)了這一點(diǎn)，因此我們決定實(shí)現(xiàn)分片多流的策略。在分片多流的過程中，相對(duì)單純的分片方式是針對(duì)單次融合的梯度進(jìn)行分片，然后多流通信。這種方式相對(duì)簡(jiǎn)單，不需要并行的處理不同的融合梯度集合，但是這種方式有一個(gè)弊端，就是在進(jìn)行下一次通信之前，必須等待本次融合切分的所有片斷都完成通信。我們進(jìn)一步發(fā)現(xiàn)，多流通信下，每條流的通信速度差異極大，因此最后使得實(shí)際速度取決于最慢的一條流，使得加速效果不明顯。因此最終采用的是完全的多梯度融合疊加多流的方式，在控制邏輯上會(huì)遠(yuǎn)比之前的邏輯復(fù)雜，整個(gè)處理過程不得不被切分為上半段和下半段。上半段負(fù)責(zé)融合以及節(jié)點(diǎn)內(nèi)Reduce再加上節(jié)點(diǎn)間通信，下半段會(huì)回到節(jié)點(diǎn)內(nèi)Broadcast結(jié)果到各節(jié)點(diǎn)。期間必須小心處理NCCL的多個(gè)流之間的關(guān)系，防止發(fā)生Hang。 延遲優(yōu)化 在深度學(xué)習(xí)的分布式訓(xùn)練中，一般情況下延遲不是主要問題。但一種情況是例外，就是在大規(guī)模的分不是訓(xùn)練下。在進(jìn)行梯度數(shù)據(jù)通信之間的梯度協(xié)商會(huì)創(chuàng)造出一個(gè)人為的單點(diǎn)熱點(diǎn)，如果按照Horovod中的處理方式，所有的節(jié)點(diǎn)均與root節(jié)點(diǎn)協(xié)商，然后root節(jié)點(diǎn)再負(fù)責(zé)協(xié)調(diào)所有節(jié)點(diǎn)。Root節(jié)點(diǎn)的負(fù)擔(dān)會(huì)隨著節(jié)點(diǎn)數(shù)增加而大幅增加，我們?cè)谝粋€(gè)客戶的CTR場(chǎng)景下，320個(gè)節(jié)點(diǎn)下，協(xié)商的overhead導(dǎo)致root節(jié)點(diǎn)完全喪失擴(kuò)展性。通過改造協(xié)商的算法，棄用了Horovod中的點(diǎn)對(duì)點(diǎn)的中心協(xié)商模式，我們將梯度協(xié)商的復(fù)雜度降低了一至兩個(gè)數(shù)量級(jí)。在CTR場(chǎng)景下也重新獲得了擴(kuò)展性。 在以上兩大類優(yōu)化之后，我們發(fā)現(xiàn)這里存在一些需要微調(diào)的參數(shù)，比如融合粒度與分流的粒度，我們最終發(fā)現(xiàn)這些參數(shù)很難離線的設(shè)定好，必須在動(dòng)態(tài)中去尋找最優(yōu)的配置。這里簡(jiǎn)要列舉一下我們做的主要的通信相關(guān)的優(yōu)化工作： 梯度融合的完全多流化，自適應(yīng)化。多梯度的融合粒度以及融合后分片的粒度，均使用自適應(yīng)的算法，在訓(xùn)練初期，從設(shè)定的一個(gè)參數(shù)空間中選擇當(dāng)前網(wǎng)絡(luò)環(huán)境下最優(yōu)的融合粒度以及分片粒度（即流數(shù)） 梯度壓縮。目前僅支持FP16的壓縮，可以根據(jù)需要，進(jìn)行scaling，來防止精度損失。后續(xù)可以將梯度壓縮插件化，允許用戶定制各種深度的梯度壓縮算法。 梯度協(xié)商的去中心化。Horovod所采用的點(diǎn)對(duì)點(diǎn)的梯度協(xié)商機(jī)制，在大節(jié)點(diǎn)時(shí)會(huì)遇到較嚴(yán)重的擴(kuò)展性問題，Perseus實(shí)現(xiàn)的去中心化的梯度協(xié)商機(jī)制有效的消除了這一問題。 性能對(duì)比 在完成了整體架構(gòu)實(shí)現(xiàn)，以及通信性能優(yōu)化后，我們對(duì)比了所有框架下的分布式性能，發(fā)現(xiàn)對(duì)大部分網(wǎng)絡(luò)模型，Horovod的性能優(yōu)于框架本身的性能，而Perseus相對(duì)于Horovod都有一定的提升。總體來說，Perseus在以下兩類狀況下，會(huì)有非常顯著的性能優(yōu)勢(shì)： 相對(duì)較大的網(wǎng)絡(luò)模型，例如NLP中的Bert，以及人臉識(shí)別中的deep insight等； 16節(jié)點(diǎn)或者更大規(guī)模的分布式訓(xùn)練； 下面是我們?cè)诳蛻糁С诌^程中，所做的一些性能測(cè)試和對(duì)比，部分網(wǎng)絡(luò)模型配合Perseus已經(jīng)運(yùn)行于客戶生產(chǎn)環(huán)境。 Tensorflow Bert Base Model  Tensorflow Bert Base Model Perseus性能為Horovod2倍至5倍 Tensorflow Resnet18  Resnet18 擬合數(shù)據(jù) Perseus 四機(jī)32卡擴(kuò)展性提升至94.2%(MXNet 為86.4%)  Resnet18真實(shí)數(shù)據(jù) Perseus 四機(jī)32卡擴(kuò)展性提升至87.4%(MXNet 為75.9%) MXNet Resnet50  Resnet50 Ali-Perseus 四機(jī)32卡擴(kuò)展性提升至91.9%(MXNet 為79.4%) MXNet Insightface  人臉識(shí)別模型 Perseus 性能為MXNet的2至2.5倍 客戶案例 通用深度學(xué)習(xí)分布式通信框架的實(shí)現(xiàn)，讓我們將各種深度學(xué)習(xí)框架的分布式優(yōu)化與框架自身完全解除了耦合，讓我們可以在一個(gè)統(tǒng)一的框架之下進(jìn)行阿里云基礎(chǔ)設(shè)施的深度優(yōu)化。我們第一階段在VPC網(wǎng)絡(luò)上的深度優(yōu)化，直接讓Tensorflow，MXNet，Pytorch和Caffe全面提升了在阿里云上得多機(jī)擴(kuò)展性，大大提升了我們平臺(tái)優(yōu)化的效率，也減輕了后期維護(hù)的工作量。目前為止，一些客戶的狀態(tài)如下： 客戶A已經(jīng)在使用Perseus+Tensorflow+Bert進(jìn)行預(yù)訓(xùn)練 客戶B在評(píng)測(cè)Perseus+MXNet，目前反饋，與客戶的自有MXNet框架整合順利，檢測(cè)模型的4機(jī)32卡性能提升10%，性能上限接近線形，客戶仍在進(jìn)一步測(cè)試。 可以看到客戶可以完全保留之前的開源訓(xùn)練框架不變的同時(shí)，享受到Perseus 通信框架的性能優(yōu)勢(shì)，對(duì)大部分客戶來說，是比較容易接受的一種方案。對(duì)于客戶購買的IaaS資源來說，會(huì)直接增加這些資源的性能，提升阿里云異構(gòu)計(jì)算產(chǎn)品的性價(jià)比。 總結(jié) Perseus的最新Release版本為0.9.4，對(duì)框架支持的情況見下圖，支持Python2.7以及3.6，支持Centos 7 以及Ubuntu 16.04。這里總結(jié)一下Perseus 通信框架幾個(gè)主要特點(diǎn)： 多訓(xùn)練框架支持 (Tensorflow,Pytorch,MXNet,Caffe) 與訓(xùn)練框架輕耦合，支持客戶的私有訓(xùn)練框架 自適應(yīng)網(wǎng)絡(luò)優(yōu)化，支持多流 梯度壓縮(支持float16) 去中心化梯度協(xié)商 NaN 檢查功能  |
|
|
