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