分布式訓練基本原理
近十年來,深度學習技術不斷刷新視覺、自然語言、語音、搜索、推薦等領域各種任務的紀錄。這其中的原因,用一個關鍵詞描述就是“大規模”。大規模的數據使得模型有足夠的知識可以記憶,大規模參數量的模型使得模型本身有能力記憶更多的數據,大規模高性能的算力(以GPU為典型代表)使得模型的訓練速度有百倍甚至千倍的提升。數據、模型、算力的發展催生了大規模深度學習這個領域,如何進行多機任務的拆分、如何配置集群訓練資源、如何平衡訓練速度和收斂速度、如何訓練單機無法訓練的模型、彈性訓練與容錯等都是這個方向重點研究的問題。分布式訓練正是解決上述問題,提升訓練效率的最有效手段。分布式訓練的核心目的是加快模型的訓練速度,如何加快訓練速度的呢?首先,可以把單機單卡的模型訓練過程看成是一個流水線生產的過程,數據就是輸入生產線的原材料,而模型的網絡可以看成生產線上的加工設備。數據經過第一個生產環節處理后,得到的結果發給下一個環節處理,后面環節依次類推。后來工廠要求提高生產效率,最直觀的方法就是再加一條生產線,分出一部分原材料交由新生產線處理,這樣生產的時間直接壓縮了一半,有哪些生產環節可以同時做的,可以直接對單個生產線進行拆分重組。分布式訓練采用了同樣的方式,即將訓練任務按照一定方法拆分到多個計算節點進行計算,再按照一定的方法對拆分后計算得到的信息,一般是梯度信息或基於梯度更新后的參數信息,進行聚合。這個拆分可以采用多種不同的方式,一般被統稱為“並行方式”,而信息聚合的方式可以稱為“模型更新傳遞方式”。下面將從這兩個方面對分布式訓練進行介紹。
分布式訓練的並行方式
在分布式訓練過程中,並不是所有的訓練任務都可以拆分,當前常用的拆分方式有如下幾種:
- 數據並行:將數據集切分放到各計算節點,每個計算節點的計算內容完全一致,並在多個計算節點之間傳遞模型參數,通常稱這種並行訓練方法為數據並行。數據並行可以解決數據集過大無法在單機高效率訓練的問題,也是工業生產中最常用的並行方法。
- 模型並行:通常指將模型單個算子計算分治到多個硬件設備上並發計算,以達到計算單個算子計算速度的目的。一般會將單個算子的計算,利用模型並行的方式分配在配置相同的幾個硬件上,進行模型存儲和計算,以保證計算步調一致。
- 流水線並行:一般是指將模型的不同算子,拆分到不同的硬件設備上進行計算,通過生產者-消費者的方式(流水線)完成不同設備之間的數據流通。深度學習中的跨設備交換數據場景很多,例如硬盤數據到內存,內存數據到顯存,內存數據到網卡等,由於不同硬件之間處理數據的速度通常不一致,通常會采用流水線並行的方式進行訓練效率最大化。
在工業場景的實踐中,分布式模型訓練也會采用不同並行方式的組合,例如數據並行與模型並行結合,數據並行與流水線並行結合等。其中數據並行的應用范圍最廣,像搜索、推薦領域的深度學習模型,通常都是采用數據並行的方式進行切分的,而后面講到的飛槳paddle分布式訓練的幾種架構方案中,基本都采用了數據並行方式,然后再與其它並行方式配合使用。
分布式訓練中的模型參數傳遞
雖然訓練任務可以拆分,但是在一個Batch的數據訓練完成后,下一個Batch的數據訓練開始之前,必須先更新Worker上的模型參數,因此模型參數的更新傳遞是分布式訓練中非常重要的一環,模型參數信息的傳遞方式和更新頻率都會直接影響模型的最終效果:
- 參數傳遞方式很大程度決定於參與分布式訓練的硬件資源選用的組織架構類型,通常有兩種架構方式,參數服務器架構和Collective架構,后面將會位大家詳細介紹。
- 參數更新頻率是指以什么樣的節奏更新參數。這同樣有兩種方式,一個是同步,即所有Worker節點,共同完成一個Batch的數據訓練后統一更新,另一個是異步,每個Worker節點,完成自己負責的數據部分的訓練后,就直接開始更新參數。
下面先來看下分布式訓練的架構。
分布式訓練架構
在分布式訓練的兩種常用架構中, Collective架構多被用於視覺、自然語言處理等需要復雜網絡計算的模型訓練任務場景,而參數服務器架構多被用於搜索推薦場景中大規模稀疏特征模型的訓練任務。下面會分別具體介紹下這兩種架構。
參數服務器架構
參數服務器架構是一種中心化架構,該架構采用將模型參數進行中心化管理的方式實現模型參數的更新和分發。參數服務器架構有兩個角色Server與Worker,值得注意的是,Server與Worker不一定對應實際的硬件,可以理解為是進程。Server負責參數的分片存儲與更新,Worker則會保存有完整的模型網絡結構,用於執行模型的前向與反向計算。常規的參數服務器的Worker節點,需要使用統一型號的CPU或GPU機器完成模型訓練。
每個Worker在執行每一步訓練時,都可以向參數服務器發出請求,獲取全局最新的模型參數,並基於最新的模型參數計算當前數據的模型梯度。當Worker基於最新的模型參數完成模型梯度的計算后,會將模型梯度發回給Server。參數服務器架構通常可以對模型參數的存儲進行分布式保存,因此對於存儲超大規模模型參數的訓練場景十分友好,這也是該架構在個性化推薦場景中,應該任務中需要保存的海量稀疏特征對應的模型參數,通常就只能采用參數服務器架構才能實現。
Collective架構
Collective架構是一種去中心化的架構,也是近年來非常流行的分布式訓練架構。該架構沒有所謂管理模型參數的中心節點,每個設備都是Worker,這個Worker同樣是進程的概念。每個Worker負責模型的訓練同時還需要掌握當前最新的全局信息。
在Collective架構中,多節點間的參數信息的傳輸通常是各個Worker之間多次點對點通信完成的,比較經典的通信算法,例如Baidu Ring All Reduce,可以采用較少的點對點通信輪數完成全局節點的模型參數傳輸。Collective架構通常在現代高性能AI芯片中使用較多,這種架構對計算芯片的算力和芯片之間的網絡互聯要求較高。高性能計算的AI芯片例如GPU,芯片之間的高速網絡互聯例如NVLINK, InfiniBand,都加快的Collective架構的發展。Collective訓練架構對於計算密集的任務非常友好,熟知的經典模型,例如機器翻譯中的Transformer,圖像分類中的ResNet50,語音識別中的DeepSpeech2通常都是采用這種訓練架構完成。
參數更新頻率
如前所述,參數同步的方式與訓練架構的選取有關,即是統一由Server更新和分發,還是Worker間相互同步,這是關系到訓練效果的重要因素之一,而另一個因素就是參數同步的頻率,即什么時候同步參數,用戶同樣可以選擇兩種方式的一種,即同步和異步,下面將會分布進行介紹。
同步訓練
同步訓練通常是指每個Worker在模型訓練的每一步都進行一次模型參數的全局更新,參數服務器架構和Collective架構都可以采用這種方式。這種方式在本質上與單機訓練一致,只是單步訓練的總數據量會更大。值得一提的是,在參數服務器架構下的同步訓練,優化算法的執行是發生在參數服務器一端,而在Collective架構下,通過對每一步所有樣本產生的模型參數梯度進行全局同步后,由每個Worker執行優化算法,相當於優化了全局的模型參數。下圖展示了單機SGD算法,基於Collective架構的同步SGD算法,基於參數服務器的異步SGD算法在計算步驟上的區別。可以看到Collective架構發生通信的位置比較集中,並行是通過與正在訓練的其他節點進行Collective通信完成,每個節點只需要關注自己即可。參數服務器架構在獲取參數,推送參數梯度等操作上都需要與參數服務器進行通信,參數服務器架構下的參數服務器端,要等待每個計算節點發送的梯度信息進行匯總后再進行模型參數的更新。
異步並行訓練算法
異步訓練在參數服務器架構下采用較多,其本質就是讓每個Worker不用關心其他節點的計算步調,獨自與參數服務器完成模型參數的更新。異步訓練情況下,參數服務器端的模型參數更新也是異步進行,即不需要等待其他正在計算的節點的進度。異步訓練下,各個計算節點的計算頻率不同,參數服務器更新模型參數的頻率也會不同,這種特性使得模型的收斂可能存在一定的問題,能夠保證高效率收斂的異步並行算法是一個研究比較多的領域。下圖是一個單機SGD與參數服務器架構下的異步SGD的計算步驟對比,可以看到參數服務器一端與同步SGD的區別在於,不需要對各個節點的梯度進行匯聚而直接進行模型參數的更新。
飛槳paddle分布式API
盡管分布式訓練能夠大大提升用戶在大規模數據下的模型訓練速度,如前面背景知識的描述可知普通用戶想快速掌握並行訓練的方法並不容易,將飛槳paddle的單機訓練程序轉換成多機訓練程序會給用戶帶來一定的使用成本。為了降低用戶使用分布式訓練的門檻,飛槳paddle官方支持分布式訓練高層API Fleet,作為分布式訓練的統一入口API,用戶可以在單機程序的基礎上進行簡單的幾行代碼修改,可實現多種類型的並行訓練方式。后面結合Fleet API介紹飛槳paddle數據並行方法,在參數服務器架構和Collective架構下,使用方法以及具體模型上的使用示例。