在過(guò)去的十年中，強(qiáng)化學(xué)習(xí)(RL)成為機(jī)器學(xué)習(xí)中最受矚目的研究領(lǐng)域之一，應(yīng)用RL能夠很好地解決芯片放置和資源管理等復(fù)雜的問(wèn)題，以及Go/Dota 2/hide-and-seek等有挑戰(zhàn)性的游戲。簡(jiǎn)單來(lái)說(shuō)，RL基礎(chǔ)架構(gòu)就是數(shù)據(jù)采集和訓(xùn)練的循環(huán)，Actor根據(jù)環(huán)境收集樣本數(shù)據(jù)，然后將其傳輸給Learner來(lái)訓(xùn)練和更新模型。當(dāng)前大多數(shù)RL實(shí)現(xiàn)都需要對(duì)環(huán)境中成千上萬(wàn)個(gè)樣本進(jìn)行多次迭代，以學(xué)習(xí)目標(biāo)任務(wù)，如Dota 2每2秒要學(xué)習(xí)成千上萬(wàn)幀樣本。這樣，RL架構(gòu)不僅要有很強(qiáng)的數(shù)據(jù)處理能力，例如增加Actor實(shí)現(xiàn)大量樣本的采集，而且還應(yīng)能夠在訓(xùn)練過(guò)程中快速迭代這些樣本。

Actor與Learner交互的RL架構(gòu)。Learner使用采樣數(shù)據(jù)訓(xùn)練模型，并將更新后的模型傳遞給Actor(例如TF-Agents，IMPALA)。

今天，我們要介紹的是Menger——一種具有本地化推理能力的大規(guī)模分布式RL架構(gòu)，可通過(guò)多個(gè)處理集群(如Borg單元)擴(kuò)展數(shù)千個(gè)Actors，從而減少了芯片放置任務(wù)的訓(xùn)練時(shí)間。在接下來(lái)的章節(jié)，我們介紹了如何使用Google TPU配置Menger，從而提高訓(xùn)練速度，進(jìn)一步我們通過(guò)芯片放置任務(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則通過(guò)分配一部分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í)間的顯著增長(zhǎng))。

在將訓(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之間引入了透明的分布式緩存組件，并通過(guò)Reverb進(jìn)行優(yōu)化(類似于Dota中使用的方法)。緩存組件的主要職責(zé)是對(duì)Actor的大量請(qǐng)求和Learner的處理能力進(jìn)行平衡。通過(guò)添加這些緩存組件，不僅顯著減輕了過(guò)多請(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)前無(wú)法擴(kuò)展到有成千上萬(wàn)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ù)載大小(16 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的分片功能來(lái)增加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í)問(wèn)題。

本文轉(zhuǎn)自雷鋒網(wǎng)，如需轉(zhuǎn)載請(qǐng)至雷鋒網(wǎng)官網(wǎng)申請(qǐng)授權(quán)。