今天，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):

# 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)以用于優化算法，如隨機梯度下降。

因此，我們提供了一種方便的方法來封裝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