0
| 本文作者： AI研習(xí)社-譯站 | 2020-10-15 11:23 |
字幕組雙語原文：Menger:大規(guī)模分布式強(qiáng)化學(xué)習(xí)架構(gòu)
英語原文：Massively Large-Scale Distributed Reinforcement Learning with Menger
在過去的十年中，強(qiáng)化學(xué)習(xí)（RL）成為機(jī)器學(xué)習(xí)中最受矚目的研究領(lǐng)域之一，應(yīng)用RL能夠很好地解決芯片放置和資源管理等復(fù)雜的問題，以及Go/Dota 2/hide-and-seek等有挑戰(zhàn)性的游戲。簡(jiǎn)單來說，RL基礎(chǔ)架構(gòu)就是數(shù)據(jù)采集和訓(xùn)練的循環(huán)，Actor根據(jù)環(huán)境收集樣本數(shù)據(jù)，然后將其傳輸給Learner來訓(xùn)練和更新模型。當(dāng)前大多數(shù)RL實(shí)現(xiàn)都需要對(duì)環(huán)境中成千上萬個(gè)樣本進(jìn)行多次迭代，以學(xué)習(xí)目標(biāo)任務(wù)，如Dota 2每2秒要學(xué)習(xí)成千上萬幀樣本。這樣，RL架構(gòu)不僅要有很強(qiáng)的數(shù)據(jù)處理能力，例如增加Actor實(shí)現(xiàn)大量樣本的采集，而且還應(yīng)能夠在訓(xùn)練過程中快速迭代這些樣本。
Actor與Learner交互的RL架構(gòu)。Learner使用采樣數(shù)據(jù)訓(xùn)練模型，并將更新后的模型傳遞給Actor（例如TF-Agents，IMPALA）。
今天，我們要介紹的是Menger——一種具有本地化推理能力的大規(guī)模分布式RL架構(gòu)，可通過多個(gè)處理集群（如Borg單元）擴(kuò)展數(shù)千個(gè)Actors，從而減少了芯片放置任務(wù)的訓(xùn)練時(shí)間。在接下來的章節(jié)，我們介紹了如何使用Google TPU配置Menger，從而提高訓(xùn)練速度，進(jìn)一步我們通過芯片放置任務(wù)驗(yàn)證框架的性能和可擴(kuò)展性?？梢园l(fā)現(xiàn)，與基準(zhǔn)模型相比，Menger將培訓(xùn)時(shí)間減少了8.6倍。
Menger設(shè)計(jì)思路
當(dāng)前有各種各樣的分布式RL系統(tǒng)，如Acme和SEED RL，然而，這些系統(tǒng)往往只從一個(gè)特定角度對(duì)分布式強(qiáng)化學(xué)習(xí)系統(tǒng)進(jìn)行優(yōu)化。例如，Acme從頻繁的Learner獲取模型，使每個(gè)Actor都進(jìn)行本地推理，而SEED RL則通過分配一部分TPU內(nèi)核執(zhí)行批量調(diào)用，進(jìn)行集中推理。對(duì)通信成本和推理成本的衡量是不同優(yōu)化系統(tǒng)的區(qū)別，具體包括：（1）向/從集中式推理服務(wù)器發(fā)送/接收觀察和動(dòng)作的通信成本，或從Larner獲取模型的通信成本；（2）相比加速器（TPU/GPU）成本，Actor的推理成本大小?？紤]到觀察值、動(dòng)作和模型大小等目標(biāo)程序要求，Menger使用類似Acme的局部推理，但同時(shí)盡可能的增加Actor的可擴(kuò)展性。要實(shí)現(xiàn)良好擴(kuò)展性和訓(xùn)練速度，主要挑戰(zhàn)包括以下兩點(diǎn)：
Actor向Learner進(jìn)行大量讀取請(qǐng)求以進(jìn)行模型檢索，這就造成Learner的負(fù)擔(dān)，隨著Actor數(shù)量的增加模型表現(xiàn)明顯受限（如收斂時(shí)間的顯著增長）。
在將訓(xùn)練數(shù)據(jù)輸送給TPU計(jì)算核心時(shí)，TPU性能通常受到輸入管道效率的限制。隨著TPU計(jì)算核心數(shù)量的增加（如TPU Pod），輸入管道的性能對(duì)于訓(xùn)練時(shí)間的影響更加明顯。
高效的模型檢索
為應(yīng)對(duì)第一個(gè)挑戰(zhàn)，在TensorFlow代碼中，我們?cè)贚earner和Actor之間引入了透明的分布式緩存組件，并通過Reverb進(jìn)行優(yōu)化（類似于Dota中使用的方法）。緩存組件的主要職責(zé)是對(duì)Actor的大量請(qǐng)求和Learner的處理能力進(jìn)行平衡。通過添加這些緩存組件，不僅顯著減輕了過多請(qǐng)求對(duì)Learner的壓力，而且以少量的通信成本將Actor分配給多個(gè)Borg單元。我們的研究表明，對(duì)有512個(gè)Actors、大小為16MB的模型，引入緩存組件可以將平均讀取延遲降低約4.0倍，從而實(shí)現(xiàn)更快的訓(xùn)練迭代，在PPO等策略算法中效果更加明顯。
若干Actors放置在不同Borg單元的分布式RL系統(tǒng)。 不同Borg單元大量Actors的頻繁的模型更新請(qǐng)求限制了Learner性能以及Learner于Actor之間的通信網(wǎng)絡(luò)，從而導(dǎo)致總體收斂時(shí)間顯著增加。虛線表示不同機(jī)器之間的gRPC通信。
引入透明分布式緩存服務(wù)的分布式RL系統(tǒng)。多個(gè)Actor放置在不同的Borg單元中，Learner僅將更新的模型發(fā)送給分布式緩存組件。每個(gè)緩存組件應(yīng)對(duì)鄰近Actor和緩存區(qū)的模型請(qǐng)求更新。緩存區(qū)不僅減輕了Learner對(duì)模型更新請(qǐng)求提供服務(wù)的負(fù)擔(dān)，而且減少了Actor的平均讀取延遲。
為提高輸入數(shù)據(jù)管道的吞吐量，Menger使用了Reverb——一種專為機(jī)器學(xué)習(xí)應(yīng)用設(shè)計(jì)的、新型開源數(shù)據(jù)存儲(chǔ)系統(tǒng)。在Reverb中，可以以在線或離線算法進(jìn)行經(jīng)驗(yàn)回放。但是，單個(gè)Reverb當(dāng)前無法擴(kuò)展到有成千上萬Actors的分布式RL系統(tǒng)，并且Actor的寫入吞吐量效率很低。
具有單個(gè)回放緩沖區(qū)的分布式RL系統(tǒng)。在Actor提出大量思維寫請(qǐng)求后，回放緩沖區(qū)會(huì)受到限制并降低總體吞吐量。此外，一旦我們將Learner擴(kuò)展到具有多個(gè)計(jì)算引擎（TPU Pod），單個(gè)回放緩沖區(qū)給這些引擎提供數(shù)據(jù)的效率就變得很低，嚴(yán)重影響總體收斂時(shí)間。
為了更好地了解回放緩沖器在分布式RL系統(tǒng)中的效率，我們?cè)u(píng)估了在不同負(fù)載大?。?6 MB-512 MB）和不同Actor（16-2048）情況下的平均寫入延遲。我們將回放緩沖區(qū)和Actor放置在同一個(gè)Borg單元中。，可以發(fā)現(xiàn)，隨著Actor數(shù)量的增加，平均寫入延遲顯著增加。將Actor的數(shù)量從16擴(kuò)展2048，16MB和512MB大小的有效負(fù)載的平均寫入延遲分別增加了約6.2倍和約18.9倍。這樣的寫入等待時(shí)間的增加影響了數(shù)據(jù)收集時(shí)間，導(dǎo)致訓(xùn)練效率低下。
將不同大小的有效負(fù)載（16 MB-512 MB）和不同數(shù)量的Actor（16至2048）放置在同一個(gè)Borg單元上時(shí)，單個(gè)Reverb重回放緩沖區(qū)的平均寫入延遲。
為緩解這種情況，我們使用Reverb的分片功能來增加Actor、Learner和和回放緩沖區(qū)之間的吞吐量。分片可在多個(gè)回放緩沖服務(wù)器之間平衡大量Actor的寫入負(fù)載，而不是僅僅作用于單個(gè)回放緩沖服務(wù)器，同時(shí)由于少數(shù)Actor共享同一服務(wù)器，可以最小化每個(gè)回放緩沖服務(wù)器的平均寫入延遲。這樣Menger就可以在多個(gè)Borg單元中擴(kuò)展數(shù)千個(gè)Actor。
具有分片回放緩沖區(qū)的分布式RL系統(tǒng)。每個(gè)回放緩沖區(qū)用于存儲(chǔ)位于同一Borg單元上的特定Actor。此外，分片回放緩沖區(qū)為加速器內(nèi)核提供了具有更高吞吐量的輸入管道。
實(shí)例驗(yàn)證：芯片放置我們面向大型網(wǎng)表在芯片放置任務(wù)中測(cè)試了Menger。與基準(zhǔn)相比，Menger使用512個(gè)TPU內(nèi)核，在訓(xùn)練時(shí)間上有了顯著改善（最高提升約8.6倍，即在最優(yōu)配置下，可以將訓(xùn)練時(shí)間從約8.6小時(shí)減少到1小時(shí)）。盡管Menger針對(duì)TPU進(jìn)行了優(yōu)化，但該框架才是性能提升的關(guān)鍵因素，我們預(yù)計(jì)在GPU上實(shí)驗(yàn)也會(huì)看到類似的提升。
與芯片放置的基準(zhǔn)相比，在不同TPU核心數(shù)量時(shí)使用Menger的訓(xùn)練時(shí)間的提升。
我們認(rèn)為，Menger架構(gòu)及其在芯片放置任務(wù)中優(yōu)異的表現(xiàn)為進(jìn)一步縮短芯片設(shè)計(jì)周期提供了方向。同時(shí)，我們還可以應(yīng)用該架構(gòu)實(shí)現(xiàn)其他具有挑戰(zhàn)性的現(xiàn)實(shí)問題。
本文工作主要有Amir Yazdanbakhsh，Chenchao Chen和Yu Zheng完成。感謝Robert Ormandi，Ebrahim Songhori，Shen Wang，TF-Agents team，Albin Cassirer，Aviral Kumar，James Laudon，John Wilkes，Joe Jiang，Milad Hashemi，Sat Chatterjee，Piotr Stanczyk，Sabela Ramos，Lasse Espeholt， Marcin Michalski，Sam Fishman，Ruoxin Sang，Azalia Mirhosseini，Anna Goldie和Eric Johnson的幫助和支持。
雷鋒字幕組是一個(gè)由AI愛好者組成的翻譯團(tuán)隊(duì)，匯聚五五多位志愿者的力量，分享最新的海外AI資訊，交流關(guān)于人工智能技術(shù)領(lǐng)域的行業(yè)轉(zhuǎn)變與技術(shù)創(chuàng)新的見解。
團(tuán)隊(duì)成員有大數(shù)據(jù)專家，算法工程師，圖像處理工程師，產(chǎn)品經(jīng)理，產(chǎn)品運(yùn)營，IT咨詢?nèi)?，在校師生；志愿者們來自IBM，AVL，Adobe，阿里，百度等知名企業(yè)，北大，清華，港大，中科院，南卡羅萊納大學(xué)，早稻田大學(xué)等海內(nèi)外高校研究所。
如果，你也是位熱愛分享的AI愛好者。歡迎與雷鋒字幕組一起，學(xué)習(xí)新知，分享成長。
雷峰網(wǎng)版權(quán)文章，未經(jīng)授權(quán)禁止轉(zhuǎn)載。詳情見轉(zhuǎn)載須知。
