本文將詳細解讀微軟與OpenAI合作推出的ZeRO-2框架、新的ONNX Runtime和超級計算機。
使用海量數據訓練模型，已成爲現代深度學習應用的一種常態。以Transformers及其它預訓練模型爲代表的一些最新的創新性應用，需要特大型訓練框架的支持。此類預訓練模型通常被視爲人工智能領域的最新突破，由此激發了全球範圍引領該領域的一些頂級技術企業間的競爭。在最近的Build大會上，微軟展示了自身在支持大規模構建和使用大規模預訓練模型方面的一些最新舉措。
深度學習領域的最新進展，主要圍繞着預訓練模型和Transformers。這場運動源自於Google提出的著名BERT模型，OpenAI GPT-2、微軟Turin-NLG等特大型架構隨後迅速跟進。建立此類模型，不僅需要佔用大量計算資源，而且通常還需要在架構上針對大規模訓練操作做出優化。由於其新穎性，主流深度學習框架尚未包含此類架構。微軟的最新舉措，全面代表了深度學習社區在簡化創建和訪問大規模模型上的進展。下面擇其重點進行概述。
Zero-2和DeepSpeed
微軟於今年年初開源了DeepSpeed的首個版本。DeepSpeed是一種優化大規模訓練的框架，是支持微軟Turing-NLG模型打破記錄的後臺基礎框架。在發佈版中，包括了一個稱爲“ZeRO”（零冗餘優化器，Zero Redundacy Optimizer）的優化模塊。從概念上講，ZeRO尋求模型並行和數據並行的中間點，致力於最大化模型的可擴展性。ZeRO的最大優點，是無需重構代碼即可支持在內存中運行大型深度學習模型。這是通過優化模型內存的使用方式而實現的。就此而言，ZeRO可視爲一種深度學習模型的內存優化方法。
作爲ZeRO的第二版，ZeRO-2相比第一版做了多處創新。其中涉及深度學習模型內存使用的多個方面，包括激活內存、碎片內存和模型狀態內存。
-
模型狀態內存（Model State Memory）： ZeRO-2優化了保存模型狀態的內存，達到相比傳統的數據並行方法降低內存使用8倍。所有深度學習模型的狀態，可歸爲分區優化器狀態、梯度和參數這三個基本過程。在ZeRO第一版中，優化了分區優化器狀態的內存使用，達到比傳統方法降低內存使用4倍。ZeRO-2在該方法上做了擴展，優化了梯度計算過程的內存佔用，進一步降低內存使用2倍。
-
激活內存（Activation Memory）：即便優化了模型狀態內存，激活函數依然可導致瓶頸。激活函數計算位於前向過程中，用於支持後向過程。ZeRO-2引入了一種將計算過程加載到主機CPU的激活分區技術，消除了內存瓶頸隱患。
-
碎片內存（Fragmented Memory）：在一些情況下，深度學習模型的低效是由於碎片所導致的，這些碎片的生成，源自於模型中各個張量的不同生命週期。ZeRO-2通過管理不同張量生命週期的內存使用全過程，消除了低效問題。
不同於其它方法中大多通過將訓練分佈在各個GPU上進行優化，ZeRO-2的一個最顯著創新就是支持單個GPU粒度的內存優化。ZeRO-2團隊展示了使用單個GPU達到在44分鐘內訓練Google BERT版本，明確給出該優化的價值所在。下圖是ZeRO-2對比以往方法的改進情況。
此外，如下圖所示，在支持具有數十億參數模型的能力、實現更高通量和高速度、每個GPU和跨GPU集羣上的擴展性、單個GPU上訓練超大模型等方面，ZeRO-2同樣展現了其優勢所在。
ZeRO-2 + ONNX Runtime
ONNX Runtime是一種支持不同機器學習框架間互操作的開源框架。機器學習互操作性的一個關注點，在於其是否適用於超大模型。ONNX可直接用於決策樹，但其是否同樣可用於具有數十億參數的模型？
爲優化ONMX模型的訓練和內存使用，微軟決定在ONNX Runtime中添加DeepSpeed和ZeRO-2。下圖展示了在一些當前全球最大的預訓練模型上的測試結果，包括在Visual Studio中的GPT-2、在Office 365中的Turing-NLG，以及Bing中具有5億參數的大型transformer。
AI優先的超級計算機
基礎架構是實現大規模預訓練模型的關鍵。微軟和OpenAI的科研人員通過合作，開發了專用於大型深度學習工作負載的超級計算機。該超級計算機在單個系統的基礎架構內，不可思議地集成了超過28.5萬個CPU內核、1萬個GPU，每臺GPU服務器的網絡互連能力達每秒400千兆位。
這臺微軟-OpenAI合作超級計算機所使用的基礎設施，已躋身全球超級計算機排名的前五位。更值得關注的是，這臺超級計算機針對AI工作負載做了優化。從這個意義上講，該超級計算機設計用於執行Turing-NLG、BERT等大型預訓練模型。微軟將新的超級計算機以Azure雲的服務形式提供給開發人員。
總 結
未來幾年中，人工智能市場最值得關注的一個趨勢，就是大規模人工智能主導地位之爭。現在DeepSpeed-ZeRO-2、ONNX Runtime和新型OpenAI超級計算機等舉措，表明微軟在該競賽中領先邁出了堅實的一步。
原文鏈接：
https://medium.com/swlh/inside-microsofts-new-frameworks-to-enable-large-scale-ai-953e9a977912