你的模型可以輕鬆使用TPU了！DeepMind 開源分散式機器學習庫TF-Replicator

新聞 03-08

新智元報道

來源：DeepMind

編輯：肖琴

【新智元導讀】今天，DeepMind又放福利：開源了一個內部的分散式機器學習庫TF-Replicator，可以幫助研究人員將TensorFlow模型輕鬆部署到GPU、TPU，並實現不同類型加速器之間的無縫切換。

最近 AI 領域的突破，從 AlphaFold 到 BigGAN 再到 AlphaStar，一個反覆出現的主題是，對方便、可靠的可擴展性的需求。

研究人員已經能夠獲取越來越多計算能力，得以訓練更大的神經網路，然而，將模型擴展到多個設備並不是一件容易的事情。

今天，DeepMind 又將其內部一個秘密武器公之於眾 ——TF-Replicator，一個可以幫助研究人員將他們的 TensorFlow 模型輕鬆部署到 GPU、Cloud TPU 的分散式機器學習框架，即使他們之前完全沒有使用分散式系統的經驗。

TF-Replicator 由 DeepMind 的研究平台團隊開發，初衷是為 DeepMind 的研究人員提供一個簡單的接入 TPU 的 API，現在，TF-Replicator 已經是DeepMind 內部最廣泛使用的 TPU 編程介面。

TF-Replicator 允許研究人員針對機器學習定位不同的硬體加速器進行，將工作負載擴展到許多設備，並在不同類型的加速器之間無縫切換。

雖然它最初是作為TensorFlow上面的一個庫開發的，但目前TF-Replicator的API已經集成到TensorFlow 2.0新的tf.distribute.Strategy中，作為 tf.distribute.Strategy的一部分開源：

https://www.tensorflow.org/alpha/guide/distribute_strategy

團隊還公開了相關論文：TF-Replicator: Distributed Machine Learning for Researchers，全面描述了這個新框架的技術細節。

https://arxiv.org/abs/1902.00465

接下來，我們將介紹 TF-Replicator 背後的想法和技術挑戰。

構建一個分散式機器學習庫

雖然 TensorFlow 為 CPU、GPU 和 TPU 設備都提供了直接支持，但是在目標之間切換需要用戶付出大量的努力。這通常涉及為特定的硬體目標專門編寫代碼，將研究想法限制在該平台的功能上。

一些構建在 TensorFlow 之上的現有框架，例如Estimators，已經試圖解決這個問題。然而，它們通常針對生產用例，缺乏快速迭代研究思路所需的表達性和靈活性。

我們開發 TF-Replicator 的初衷是為 DeepMind 的研究人員提供一個使用 TPU 的簡單API。TPU 為機器學習工作負載提供了可擴展性，實現了許多研究突破，例如使用我們的 BigGAN 模型實現了最先進的圖像合成。

TensorFlow 針對 TPU 的原生 API 與針對 GPU 的方式不同，這造成了使用 TPU 的障礙。TF-Replicator 提供了一個更簡單、更用戶友好的 API，隱藏了 TensorFlow 的 TPU API 的複雜性。此外，研究平台團隊與不同機器學習領域的研究人員密切合作，開發了TF-Replicator API，以確保必要的靈活性和易用性。

TF-Replicator API

使用 TF-Replicator 編寫的代碼與使用 TensorFlow 中為單個設備編寫的代碼類似，允許用戶自由定義自己的模型運行循環。用戶只需要定義 (1) 一個公開數據集的輸入函數，以及 (2) 一個定義其模型邏輯的 step 函數 (例如，梯度下降的單個 step):

將計算擴展到多個設備需要設備之間進行通信。在訓練機器學習模型的背景下，最常見的通信形式是累積梯度(accumulate gradients) 以用於優化演算法，如隨機梯度下降。

因此，我們提供了一種方便的方法來封裝 TensorFlow Optimizers，以便在更新模型參數之前在設備之間累積梯度。對於更一般的通信模式，我們提供了類似於 MPI 的原語，如「all_reduce」和「broadcast」。這些使得實現諸如全局批標準化之類的操作變得非常簡單，這是擴展 BigGAN 模型訓練的關鍵技術。

輸入數據從主機發送到各個 GPU, GPU 立即開始處理。當需要在 GPU 之間交換信息時，它們會在發送數據之前進行同步。

實現

對於多 GPU 計算，TF-Replicator 依賴於「圖內複製」(「in-graph replication)模式，其中每個設備的計算在同一個 TensorFlow graph 中複製。設備之間的通信是通過連接設備對應子圖中的節點來實現的。在 TF-Replicator 中實現這一點很具挑戰性，因為在TensorFlow graph 中的任何位置都可能發生通信。因此，構造計算的順序至關重要。

我們的第一個想法是在一個單獨的 Python 線程中同時構建每個設備的子圖。當遇到通信原語時，線程同步，主線程插入所需的跨設備計算。之後，每個線程將繼續構建其設備的計算。

然而，在我們考慮這種方法時，TensorFlow 的圖形構建 API 不是線程安全的，這使得在不同線程中同時構建子圖非常困難。相反，我們使用圖形重寫 (graph rewriting) 在所有設備的子圖構建完成後插入通信。在構造子圖時，佔位符被插入到需要通信的位置。然後，我們跨設備收集所有匹配佔位符，並用適當的跨設備計算替換它們。

當 TF-Replicator 構建一個 in-graph replicated 計算時，它首先獨立地為每個設備構建計算，並將佔位符留給用戶指定的跨設備計算。構建好所有設備的子圖之後，TF-Replicator 通過用實際的跨設備計算替換佔位符來連接它們。

為 AI 研究構建一個平台

通過在 TF-Replicator 的設計和實現過程中與研究人員密切合作，我們最終構建一個庫，讓用戶能夠輕鬆地跨多個硬體加速器進行大規模計算，同時讓他們擁有進行前沿 AI研究所需的控制和靈活性。

例如，在與研究人員討論之後，我們添加了 MPI 風格的通信原語，如 all-reduce。TF-Replicator 和其他共享基礎架構使我們能夠在穩健的基礎上構建越來越複雜的實驗，並在整個 DeepMind 快速傳播最佳實踐。

在撰寫本文時，TF-Replicator 已經成為 DeepMind 應用最廣泛的 TPU 編程介面。雖然這個庫本身並不局限於訓練神經網路，但它最常用來訓練大量數據。例如，BigGAN 模型是在一個 512 核的 TPUv3 pod 訓練的，batch size 為 2048。

在採用分散式 actor-learner 設置的增強學習智能體中，例如我們的重要性加權 actor-learner 架構，可擴展性是通過讓許多 actor 通過與環境的交互生成新的體驗來實現的。然後，learner 對這些數據進行處理，以改進 agent 的策略，表示為一個神經網路。為了應對越來越多的actor，TF-Replicator 可以很輕鬆地將 learner 分布在多個硬體加速器上。

這些以及更多例子在我們的 arXiv 論文中有更詳細的描述。

Blog:

https://deepmind.com/blog/tf-replicator-distributed-machine-learning/

Paper：

https://arxiv.org/abs/1902.00465

更多閱讀

【加入社群】

新智元AI技術產業社群招募中，歡迎對AI技術產業落地感興趣的同學，加小助手微信號：aiera2015_2入群;通過審核後我們將邀請進群，加入社群後務必修改群備註（姓名 - 公司 - 職位;專業群審核較嚴，敬請諒解）