|
雷鋒網(wǎng) AI 科技評論按：DeepMind 有一支專門的科研平臺團(tuán)隊（the Research Platform Team），他們的職責(zé)是為 AI 學(xué)術(shù)研究構(gòu)建加速計算的基礎(chǔ)設(shè)施。他們不經(jīng)常亮相，但是這次由他們撰文介紹的 TF-Replicator 是一個極為有用的工具：它是又一個實用的軟件庫，可以幫助從未接觸過分布式系統(tǒng)的研究人員們輕松地在 GPU 集群和云 TPU 集群上部署 TensorFlow 模型，也就成為了深度學(xué)習(xí)進(jìn)行大規(guī)模工業(yè)化應(yīng)用的重要組件。TF-Replicator 的程序化模型現(xiàn)在也已經(jīng)作為 TensorFlow 的 tf.distribute.Strategy 的一部分開源在 https://www./alpha/guide/distribute_strategy。 DeepMind 科研平臺團(tuán)隊發(fā)表的技術(shù)博客中介紹了 TF-Replicator 的核心思想、開發(fā)經(jīng)歷、技術(shù)挑戰(zhàn)。雷鋒網(wǎng) AI 科技評論把主要內(nèi)容編譯如下。 在 AlphaFold、BigGAN、AlphaStar 等近期的人工智能科研成果中，我們很容易看到一個反復(fù)出現(xiàn)的要素，那就是對輕松增加模型和計算能力規(guī)模的渴求。在越來越高的計算能力上訓(xùn)練越來越大的模型，讓神經(jīng)網(wǎng)站展現(xiàn)出了驚人的表現(xiàn)。為了應(yīng)對這樣的需求，DeepMind 科研平臺團(tuán)隊開發(fā)了 TF-Replicator，它可以幫助科研人員輕松地為機(jī)器學(xué)習(xí)模型指定不同的硬件加速器、把運(yùn)算拓展到更多設(shè)備上，以及在不同類型的加速器之間無縫切換。TF-Replicator 一開始是基于 TensorFlow 進(jìn)行開發(fā)的，如今它的 API 也已經(jīng)直接集成進(jìn)了 TF 2.0 中。 TensorFlow 固然對 CPU、GPU、TPU 都提供了直接的支持，但是用戶想要切換模型運(yùn)行在哪個設(shè)備上卻是一件很麻煩的事情。一般來說，這個過程都需要針對目標(biāo)硬件重新編寫專門的代碼，這也就限制了科研想法在具體平臺上的發(fā)揮。目前也有一些別的基于 TensorFlow 的框架在嘗試應(yīng)對這樣的問題，比如 Estimators，不過這些方案基本上都是針對生產(chǎn)環(huán)境的，缺少學(xué)術(shù)科研的快速迭代所需要的表達(dá)能力和靈活性。 構(gòu)建一個分布式機(jī)器學(xué)習(xí)庫DeepMind 科研平臺團(tuán)隊開發(fā) TF-Replicator 的初始目標(biāo)是為 DeepMind 的研究人員們提供一種簡單便捷的使用 TPU 的 API。不過 TensorFlow 原生的 TPU 計算 API 和 GPU 計算 API 有不少區(qū)別，這就為切換到 TPU 帶來了不小障礙。TF-Replicator 提供的 API 要更簡單、更友好，隱藏了 TF 原本的 TPU API 中的復(fù)雜之處。其中最關(guān)鍵的是，科研平臺團(tuán)隊和 DeepMind 之外的許多機(jī)器學(xué)習(xí)組織機(jī)構(gòu)都進(jìn)行了緊密的合作，以確保他們開發(fā)出的 TF-Replicator 有足夠的靈活性、易用性。 TF-Replicator API使用了 TF-Replicator 的代碼看起來和單臺設(shè)備上運(yùn)行的 TF 代碼非常相似，給了用戶足夠的自由來定義他們自己的模型運(yùn)行循環(huán)。用戶只需要做兩次簡單的定義，一，定義指向某個數(shù)據(jù)集的輸入函數(shù)；二，定義一個指定了模型邏輯的步驟函數(shù)（比如梯度下降中的某一步迭代）。代碼示例如下：
想要把運(yùn)算拓展到多個設(shè)備，需要讓這些設(shè)備之間能夠相互溝通。在訓(xùn)練機(jī)器學(xué)習(xí)模型的情境中，最常見的溝通形式就是為隨機(jī)梯度下降 SGD 之類的優(yōu)化算法累積計算梯度。所以 DeepMind 科研平臺團(tuán)隊在 TF-Replicator 中用一個便捷的方法提供了 TensorFlow 各種優(yōu)化器的打包，這樣在更新模型的參數(shù)之前，不同設(shè)備分別計算的梯度就可以先進(jìn)行累積。對于其它的更通用化的溝通，TF-Replicator 提供了類似 MPI 的主動式溝通模式，比如 `all_reduce` 和 `broadcast`。這些功能的幫助下，BigGAN 模型中使用到的全局批量歸一化（global batch normalisation）也可以簡單地在操作中實現(xiàn)，而這也是 BigGAN 增加訓(xùn)練規(guī)模非常重要的一步（可以參見 https:///abs/1809.11096 論文第三節(jié)）。 對于多 GPU 配置中的運(yùn)算，TF-Replicator 使用的是一種“圖內(nèi)復(fù)制”（in-graph replication）的模式，每個設(shè)備的計算都會復(fù)制到同一個 TensorFlow 圖中。設(shè)備之間的溝通方式是連接不同設(shè)備的對應(yīng)子圖的節(jié)點。在 TF-Replicator 中實現(xiàn)這個其實很有挑戰(zhàn)性，因為溝通可以發(fā)生在數(shù)據(jù)流圖中的任何一點上。運(yùn)算的構(gòu)建順序就變得尤其重要。 DeepMind 科研平臺團(tuán)隊最開始的想法是在單獨的 Python 線程中為每個設(shè)備同步構(gòu)建各自的子圖。當(dāng)遇到一個主動發(fā)起溝通的請求時，這些線程就會同步，主線程會插入執(zhí)行所需的跨設(shè)備運(yùn)算。在此之后，每個線程都會繼續(xù)構(gòu)建它自己對應(yīng)的設(shè)備的運(yùn)算。然而，對這種方法進(jìn)行仔細(xì)思考之后他們意識到 TF 中的圖構(gòu)建 API 并不是線程安全的，這意味著同步在不同的線程中構(gòu)建不同的子圖非常困難。所以他們轉(zhuǎn)而使用了圖重寫（graph rewriting），在所有設(shè)備的子圖都完成構(gòu)建之后插入溝通。當(dāng)這些子圖還在構(gòu)建的時候，TF-Replicator 會在需要溝通的地方插入一個占位符，子圖構(gòu)建完成之后它就會在不同的設(shè)備間比對所有匹配的占位符，然后把它們更換成正確的跨設(shè)備運(yùn)算。 由于在設(shè)計和實現(xiàn)的全過程中都和學(xué)術(shù)界內(nèi)的研究人員們保持緊密的合作，TF-Replicator 得以被構(gòu)建成一個可以輕松地在多種硬件加速器之間拓展運(yùn)算的庫，同時還讓用戶保留了前沿 AI 研究所需的高度控制和靈活性。 在開發(fā)完畢之前，TF-Replicator 就已經(jīng)是 DeepMind 的研究人員們使用最多的 TPU 編程接口。TF-Replicator 的功能并不限制于訓(xùn)練神經(jīng)網(wǎng)絡(luò)，不過它還是最常用來在大量數(shù)據(jù)上進(jìn)行大批量訓(xùn)練。比如 BigGAN 模型就是在最高達(dá)到 512 個 TPUv3 核心組成的集群上，以 2048 的批量大小進(jìn)行訓(xùn)練的。在帶有分布式的執(zhí)行者-學(xué)習(xí)者設(shè)置的強(qiáng)化學(xué)習(xí)智能體中，更大規(guī)模的運(yùn)算可以讓更多不同的執(zhí)行者和環(huán)境進(jìn)行交互，也可以產(chǎn)生新的經(jīng)驗。TF-Replicator 已經(jīng)為更多執(zhí)行者的環(huán)境做好了準(zhǔn)備，可以輕松地在多個硬件加速器上分布學(xué)習(xí)者（通常是神經(jīng)網(wǎng)絡(luò)）的運(yùn)算。在論文中有更詳細(xì)的介紹。 DeepMind 科研平臺團(tuán)隊開發(fā)了許多有影響力的技術(shù)工具，TF-Replicator 只是其中之一。DeepMind 的許多深度學(xué)習(xí)科研成果，都有科研平臺團(tuán)隊的功勞。DeepMind 科研平臺團(tuán)隊在介紹分享自己的成果的同時，也歡迎對前沿 AI 研究感興趣的科研平臺開源軟件工程師加入他們。 詳細(xì)介紹請見論文 https:///abs/1902.00465 via DeepMind Blog，雷鋒網(wǎng) AI 科技評論編譯 |
|
|
