導讀:下一代人工智能應用程序需要不斷地與環境交互,并從這些交互中學習。這對系統的性能和靈活性提出了新的要求,而現有的機器學習計算框架大多無法滿足這些要求。為此,UC Berkeley 項目組開發了一個新的分布式框架 Ray,并于近日在 Arvix 上發表了相關論文:《Ray: A Distributed Framework for Emerging AI Applications》。
論文第一作者為 Philipp Moritz 及 Robert Nishihara,是 UC Berkeley AMP Lab 的博士生,而 Michael I. Jordan 和 Ion Stoica 的名字也赫然列于其中。
Michael I. Jordan :UC Berkeley 電氣工程與計算機科學系和統計系杰出教授,是美國國家科學院、美國國家工程院、美國藝術與科學院三院院士,是機器學習領域唯一獲此成就的科學家。2016 年,他被 Semantic Scholar 評為“最具影響力的計算機科學家”。
Ion Stoica :UC Berkeley 計算機系教授,AMPLab 共同創始人,彈性 P2P 協議 Chord、集群內存計算框架 Spark、集群資源管理平臺 Mesos 核心作者。
目前的計算框架存在的短板
如今大部分人工智能應用都是基于局限性較大的監督學習的范式而開發的,即模型在線下進行訓練,然后部署到服務器上進行線上預測。隨著該領域的成熟,機器學習應用需要更多地在動態環境下運行,響應環境中的變化,并且采用一系列的動作來完成既定目標。這些要求自然地建立在增強學習(Reinforcement Learning,RL)范式中,即在不確定的環境中連續學習。
RL 應用與傳統的監督學習應用有三個不同之處:
1)RL 應用嚴重依賴仿真來探索所在狀態及操作結果。這需要大量的計算,現實情況下,一個應用大概需要進行億萬次仿真。
2)RL 應用的計算圖是異質的、動態變化的。一次仿真可能會花掉幾毫秒到幾分鐘的時間,仿真的結果又決定未來仿真的參數。
3)許多 RL 應用程序,如機器人控制或自主駕駛,需要迅速采取行動,以響應不斷變化的環境。
因此,我們需要一個能支持異質和動態計算圖,同時以毫秒級延遲每秒處理數以百萬計任務的計算框架。而目前的計算框架或是無法達到普通 RL 應用的延遲要求(MapReduce、Apache Spark、CIEL),或是使用靜態計算圖(TensorFlow、Naiad、MPI、Canary)。
RL 應用對系統提出了靈活性、表現性能以及易開發的要求,Ray 系統則是為滿足這些要求而設計的。
示例
經典RL 訓練應用偽代碼
用Ray 實現的Python 代碼樣例
在Ray 中,通過@ray.remote 聲明remote 函數和actor。當調用remote 函數和actor methods 時會立即返回一個future(對象id),使用ray.get()可以同步獲取該id 對應的對象,可以傳遞給后續的remote 函數和actor methods 來編碼任務依賴項。每個actor 有一個環境對象 self.env,在任務之間共享狀態。
上圖是調用train_policy.remote() 對應的任務圖。remote 函數和actor methods 調用對應任務圖中的任務。圖中有2 個actor,每個actor 之間的狀態邊(stateful edges)意味著他們共享可變狀態。從train_policy 到它所調用的任務之間有控制邊(control edges)。要并行訓練策略(policy),可以多次調用train_policy.remote()。
原理
為了支持RL 應用所帶來的異質和動態工作負荷要求,Ray 采用與CIEL 類似的動態任務圖計算模型。除了CIEL 的任務并行簡化外,Ray 在執行模型頂層提供了代碼簡化,能夠支持諸如第三方仿真的狀態結構。
Ray 系統結構
為了在支持動態計算圖的同時滿足嚴格的性能要求,Ray 采取一種新的可橫向擴展的分布式結構。Ray 的結構由兩部分組成:application 層和 system 層。Application 層實現 API 和計算模型,執行分布式計算任務。System 層負責任務調度和數據管理,來滿足表現性能和容錯的要求。
Ray 系統結構
該結構基于兩個關鍵想法:
1)全局狀態存儲 GSC(Global Control Store)。系統所有的控制狀態存儲在 GSC 中,這樣系統其他組件可以是無狀態的。不僅簡化了對容錯的支持(出現錯誤時,組件可以從 GSC 中讀取最近狀態并重新啟動),也使得其他組件可以橫向擴展(該組件的復制或碎片可以通過 GSC 狀態共享)。
2)自底向上的分布式調度器。任務由 driver 和 worker 自底向上地提交給局部調度器(local scheduler)。局部調度器可以選擇局部調度任務,或將任務傳遞給全局調度器。通過允許本地決策,降低了任務延遲,并且通過減少全局調度器的負擔,增加了系統的吞吐量。
自底向上的分布式調度器
性能表現
1)可擴展性和表現性能
端到端可擴展性。 GCS 的主要優勢是增強系統的橫向可擴展性。我們可以觀察到幾乎線性的任務吞吐量增長。在 60 節點,Ray 可以達到超過每秒 100 萬個任務的吞吐量,并線性地在 100 個節點上超過每秒 180 萬個任務。最右邊的數據點顯示,Ray 可以在不到一分鐘的時間處理 1 億個任務(54s)。
全局調度器的主要職責是在整個系統中保持負載平衡。Driver 在第一個節點提交了100K 任務,由全局調度器平衡分配給21 個可用節點。
對象存儲性能。對于大對象,單一客戶端吞吐量超過了15GB/s(紅色),對于小對象,對象存儲IOPS 達到18K(青色),每次操作時間約56 微秒。
2)容錯性
從對象失敗中恢復。隨著 worker 節點被終結,活躍的局部調度器會自動觸發丟失對象重建。在重建期間,driver 最初提交的任務被擱置,因為它們的依賴關系不能滿足。但是整體的任務吞吐量保持穩定,完全利用可用資源,直到丟失的依賴項被重建。
分布式任務的完全透明容錯。虛線表示集群中的節點數。曲線顯示新任務(青色)和重新執行任務(紅色)的吞吐量,到210s 時,越來越多的節點加回到系統,Ray 可以完全恢復到初始的任務吞吐量。
從actor 失敗中恢復。通過將每個actor 的方法調用編碼到依賴關系圖中,我們可以重用同一對象重構機制。
t=200s 時,我們停止 10 個節點中的 2 個,導致集群中 2000 個 actor 中的 400 個需要在剩余節點上恢復。(a)顯示的是沒有中間節點狀態被存儲的極端情況。調用丟失的 actor 的方法必須重新串行執行(t = 210-330s)。丟失的角色將自動分布在可用節點上,吞吐量在重建后完全恢復。(b)顯示的是同樣工作負載下,每 10 次方法調用每個 actor 自動進行了一次 checkpoint 存儲。節點失效后,大部分重建是通過執行 checkpoint 任務重建 actor 的狀態(t = 210-270s)。
GCS 復制消耗。為了使 GCS 容錯,我們復制每個數據庫碎片。當客戶端寫入 GCS 的一個碎片時,它將寫入復制到所有副本。通過減少 GCS 的碎片數量,我們人為地使 GCS 成為工作負載的瓶頸,雙向復制的開銷小于 10%。
3)RL 應用
我們用 Ray 實現了兩種 RL 算法,與專為這兩種算法設計的系統進行對比,Ray 可以趕上甚至超越特定的系統。除此之外,使用 Ray 在集群上分布這些算法只需要在算法實現中修改很少幾行代碼。
ES 算法(Evolution Strategies)
Ray 和參考系統實現 ES 算法在 Humanoid v1 任務上達到 6000 分所需時間對比。
在 Ray 上實現的 ES 算法可以很好地擴展到 8192 核,而特制的系統在 1024 核后便無法運行。在 8192 核上,我們取得了中值為 3.7 分鐘的效果,比目前最好效果快兩倍。
PPO 算法(Proximal Policy Optimization)
為了評估 Ray 在單一節點和更小 RL 工作負載的性能,我們在 Ray 上實現了 PPO 算法,與 OpenMPI 實現的算法進行對比。
MPI 和 Ray 實現 PPO 算法在 Humanoid v1 任務上達到 6000 分所需時間對比。
用 Ray 實現的 PPO 算法超越了特殊的 MPI 實現,并且使用 GPU 更少。
控制仿真機器人
實驗表明,Ray 可以達到實時控制模擬機器人的軟實時要求。Ray 的驅動程序能運行模擬機器人,并在固定的時間間隔采取行動,從 1 毫秒到 30 毫秒,以模擬不同的實時要求。
未來工作
考慮到工作負載的普遍性,特殊的優化是比較難的。例如,必須在沒有計算圖的全部知識情況下采取調度決策。Ray 的調度決策或許需要更復雜的設置。除此之外,每個任務的存儲譜系需要執行垃圾收集策略,以在 GCS 中限制存儲成本,這是目前正在開發的功能。
當 GCS 的消耗成為瓶頸時,可以通過增加更多的碎片來擴展全局調度器。目前還需要手動設置 GCS 碎片和全局調度器的數量,未來將開發自適應算法來自動調整它們的數量??紤]到 GCS 結構為該系統帶來的優勢,作者認為集中化控制狀態是未來分布式系統的關鍵設計元素。