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

最近AI領(lǐng)域的突破，從AlphaFold到BigGAN再到AlphaStar，一個(gè)反復(fù)出現(xiàn)的主題是，對(duì)方便、可靠的可擴(kuò)展性的需求。

研究人員已經(jīng)能夠獲取越來越多計(jì)算能力，得以訓(xùn)練更大的神經(jīng)網(wǎng)絡(luò)，然而，將模型擴(kuò)展到多個(gè)設(shè)備并不是一件容易的事情。

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

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

TF-Replicator允許研究人員針對(duì)機(jī)器學(xué)習(xí)定位不同的硬件加速器進(jìn)行，將工作負(fù)載擴(kuò)展到許多設(shè)備，并在不同類型的加速器之間無縫切換。

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

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

團(tuán)隊(duì)還公開了相關(guān)論文：TF-Replicator: Distributed Machine Learning for Researchers，全面描述了這個(gè)新框架的技術(shù)細(xì)節(jié)。

https://arxiv.org/abs/1902.00465

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

構(gòu)建一個(gè)分布式機(jī)器學(xué)習(xí)庫

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

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

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

TensorFlow針對(duì)TPU的原生API與針對(duì)GPU的方式不同，這造成了使用TPU的障礙。TF-Replicator提供了一個(gè)更簡(jiǎn)單、更用戶友好的API，隱藏了TensorFlow的TPU API的復(fù)雜性。此外，研究平臺(tái)團(tuán)隊(duì)與不同機(jī)器學(xué)習(xí)領(lǐng)域的研究人員密切合作，開發(fā)了TF-Replicator API，以確保必要的靈活性和易用性。

TF-Replicator API

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

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

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

輸入數(shù)據(jù)從主機(jī)發(fā)送到各個(gè)GPU, GPU立即開始處理。當(dāng)需要在GPU之間交換信息時(shí)，它們會(huì)在發(fā)送數(shù)據(jù)之前進(jìn)行同步。

實(shí)現(xiàn)

對(duì)于多GPU計(jì)算，TF-Replicator依賴于“圖內(nèi)復(fù)制”(“in-graph replication)模式，其中每個(gè)設(shè)備的計(jì)算在同一個(gè)TensorFlow graph中復(fù)制。設(shè)備之間的通信是通過連接設(shè)備對(duì)應(yīng)子圖中的節(jié)點(diǎn)來實(shí)現(xiàn)的。在TF-Replicator中實(shí)現(xiàn)這一點(diǎn)很具挑戰(zhàn)性，因?yàn)樵赥ensorFlow graph中的任何位置都可能發(fā)生通信。因此，構(gòu)造計(jì)算的順序至關(guān)重要。

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

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

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

為AI研究構(gòu)建一個(gè)平臺(tái)

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

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

在撰寫本文時(shí)，TF-Replicator已經(jīng)成為DeepMind應(yīng)用最廣泛的TPU編程接口。雖然這個(gè)庫本身并不局限于訓(xùn)練神經(jīng)網(wǎng)絡(luò)，但它最常用來訓(xùn)練大量數(shù)據(jù)。例如，BigGAN模型是在一個(gè)512核的TPUv3 pod訓(xùn)練的，batch size為2048。

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

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

Blog:

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

Paper：

https://arxiv.org/abs/1902.00465