DeepMind又放福利：開源了一個內部的分布式機器學習庫TF-Replicator

今天，DeepMind又放福利：開源了一個內部的分布式機器學習庫TF-Replicator，可以幫助研究人員將TensorFlow模型輕松部署到GPU、TPU，并實現(xiàn)不同類型加速器之間的無縫切換。

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

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

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

TF-Replicator由DeepMind的研究平臺團隊開發(fā)，初衷是為DeepMind的研究人員提供一個簡單的接入TPU的API，現(xiàn)在，TF-Replicator已經是DeepMind內部最廣泛使用的TPU編程接口。

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

雖然它最初是作為TensorFlow上面的一個庫開發(fā)的，但目前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，全面描述了這個新框架的技術細節(jié)。

https://arxiv.org/abs/1902.00465

接下來，我們將介紹TF-Replicator背后的想法和技術挑戰(zhàn)。

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

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

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

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

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

TF-Replicator API

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

# Deploying a model with TpuReplicator.repl = tf_replicator.TpuReplicator( num_workers=1, num_tpu_cores_per_worker=8)with repl.context(): model = resnet_model() base_optimizer = tf.train.AdamOptimizer() optimizer = repl.wrap_optimizer(base_optimizer)# ... code to define replica input_fn and step_fn.per_replica_loss = repl.run(step_fn, input_fn)train_op = tf.reduce_mean(per_replica_loss)with tf.train.MonitoredSession() as session: repl.init(session) for i in xrange(num_train_steps): session.run(train_op) repl.shutdown(session)

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

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

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

實現(xiàn)

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

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

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

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

為AI研究構建一個平臺

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

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

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

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

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

Blog:

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

Paper：

https://arxiv.org/abs/1902.00465