利用深度強(qiáng)化學(xué)習(xí)設(shè)計(jì)算術(shù)電路

隨著摩爾定律的放緩，開發(fā)其他技術(shù)來(lái)提高同一技術(shù)過(guò)程節(jié)點(diǎn)上芯片的性能變得越來(lái)越重要。我們的方法使用人工智能設(shè)計(jì)更小、更快、更高效的電路，以在每一代芯片中提供更高的性能。

大量的算術(shù)電路陣列為 NVIDIA GPU 提供了動(dòng)力，使其在人工智能、高性能計(jì)算和計(jì)算機(jī)圖形學(xué)方面實(shí)現(xiàn)了前所未有的加速。因此，改進(jìn)這些算術(shù)電路的設(shè)計(jì)對(duì)于提高 GPU 的性能和效率至關(guān)重要。

如果人工智能能夠?qū)W會(huì)設(shè)計(jì)這些電路呢？在 PrefixRL ：使用深度強(qiáng)化學(xué)習(xí)優(yōu)化并行前綴電路 中，我們證明了人工智能不僅可以從頭開始學(xué)習(xí)設(shè)計(jì)這些電路，而且人工智能設(shè)計(jì)的電路也比最先進(jìn)的電子設(shè)計(jì)自動(dòng)化（ EDA ）工具設(shè)計(jì)的電路更小更快。最新 NVIDIA Hopper GPU 結(jié)構(gòu) 擁有近 13000 個(gè)人工智能設(shè)計(jì)電路實(shí)例。

圖 1 。由 PrefixRL AI （左）設(shè)計(jì)的 64b 加法器電路比由最先進(jìn)的 EDA 工具（右）設(shè)計(jì)的電路小 25% ，同時(shí)速度快，功能等效

在圖 1 中，電路對(duì)應(yīng)于圖 5 中 PrefixRL 曲線中的（ 31.4 μ m 2， 0.186ns ）點(diǎn)。

電路設(shè)計(jì)游戲

計(jì)算機(jī)芯片中的算術(shù)電路是使用邏輯門（如 NAND 、 NOR 和 XOR ）和導(dǎo)線組成的網(wǎng)絡(luò)構(gòu)建的。理想電路應(yīng)具有以下特性：

Small： 一個(gè)較低的區(qū)域，以便在一個(gè)芯片上可以容納更多的電路。

Fast： 降低延遲以提高芯片性能。

耗電更少： 芯片功耗更低。

在本文中，我們重點(diǎn)研究了電路面積和延遲。我們發(fā)現(xiàn)功耗與感興趣電路的面積密切相關(guān)。電路面積和延遲通常是相互競(jìng)爭(zhēng)的屬性，因此我們希望找到有效權(quán)衡這些屬性的設(shè)計(jì)的帕累托前沿。簡(jiǎn)單地說(shuō)，我們希望在每個(gè)延遲時(shí)都有最小面積電路。

在 PrefixRL 中，我們關(guān)注一類流行的算術(shù)電路，稱為（并行）前綴電路。 GPU 中的各種重要電路，如加法器、增量器和編碼器，都是前綴電路，可以在更高級(jí)別上定義為前綴圖。

在這項(xiàng)工作中，我們特別提出了一個(gè)問題：人工智能代理能否設(shè)計(jì)出良好的前綴圖？所有前綴圖的狀態(tài)空間都很大O（2^n^n），無(wú)法使用蠻力方法進(jìn)行探索。

圖 2:PrefixRL 的一次迭代與 4b 電路示例

使用電路生成器將前綴圖轉(zhuǎn)換為具有導(dǎo)線和邏輯門的電路。然后，使用物理合成優(yōu)化（如柵極尺寸、復(fù)制和緩沖器插入）通過(guò)物理合成工具進(jìn)一步優(yōu)化這些生成的電路。

由于這些物理合成優(yōu)化，最終電路特性（延遲、面積和功率）不會(huì)直接從原始前綴圖特性（如電平和節(jié)點(diǎn)數(shù)）轉(zhuǎn)換。這就是為什么人工智能代理學(xué)習(xí)設(shè)計(jì)前綴圖，但優(yōu)化由前綴圖生成的最終電路的屬性。

我們將算術(shù)電路設(shè)計(jì)作為強(qiáng)化學(xué)習(xí)（ RL ）任務(wù)，在該任務(wù)中，我們訓(xùn)練代理優(yōu)化算術(shù)電路的面積和延遲特性。對(duì)于前綴電路，我們?cè)O(shè)計(jì)了一個(gè)環(huán)境，在該環(huán)境中， RL 代理可以在前綴圖中添加或刪除節(jié)點(diǎn)，然后執(zhí)行以下步驟：

前綴圖合法化，以始終保持正確的前綴和計(jì)算。

從合法化的前綴圖生成電路。

使用物理合成工具對(duì)電路進(jìn)行物理合成優(yōu)化。

測(cè)量了電路的面積和延遲特性。

在一集中， RL 代理通過(guò)添加或刪除節(jié)點(diǎn)逐步建立前綴圖。在每個(gè)步驟中，代理都會(huì)收到相應(yīng)電路區(qū)域的改進(jìn)和延遲作為獎(jiǎng)勵(lì)。

狀態(tài)和動(dòng)作表示與深度強(qiáng)化學(xué)習(xí)模型

我們使用 Q 學(xué)習(xí)算法來(lái)訓(xùn)練電路設(shè)計(jì)代理。我們對(duì)前綴圖使用網(wǎng)格表示，其中網(wǎng)格中的每個(gè)元素唯一地映射到前綴節(jié)點(diǎn)。這種網(wǎng)格表示法用于 Q 網(wǎng)絡(luò)的輸入和輸出。輸入網(wǎng)格中的每個(gè)元素表示節(jié)點(diǎn)是否存在。輸出網(wǎng)格中的每個(gè)元素表示用于添加或刪除節(jié)點(diǎn)的 Q 值。

我們使用完全卷積神經(jīng)網(wǎng)絡(luò)架構(gòu)作為智能體的輸入和輸出， Q 學(xué)習(xí)智能體是網(wǎng)格表示。該代理單獨(dú)預(yù)測(cè)面積和延遲屬性的 Q 值，因?yàn)槊娣e和延遲的回報(bào)在訓(xùn)練期間是可以單獨(dú)觀察到的。

圖 3 。某些 4b 前綴圖（左）和全卷積 Q 學(xué)習(xí)代理架構(gòu)（右）的表示

Raptor 分布式訓(xùn)練

PrefixRL 是一項(xiàng)計(jì)算要求很高的任務(wù)：物理模擬每個(gè) GPU 需要 256 CPU ，訓(xùn)練 64b 案例需要 32000 GPU 小時(shí)。

我們開發(fā)了 Raptor ，這是一個(gè)內(nèi)部分布式強(qiáng)化學(xué)習(xí)平臺(tái)，它利用了 NVIDIA 硬件的特殊優(yōu)勢(shì)來(lái)進(jìn)行這種工業(yè)強(qiáng)化學(xué)習(xí)（圖 4 ）。

Raptor 具有一些增強(qiáng)可擴(kuò)展性和訓(xùn)練速度的功能，例如作業(yè)調(diào)度、自定義網(wǎng)絡(luò)和 GPU 感知的數(shù)據(jù)結(jié)構(gòu)。在 PrefixRL 的上下文中， Raptor 使工作分布在 CPU 、 GPU 和 Spot 實(shí)例的混合中成為可能。

此強(qiáng)化學(xué)習(xí)應(yīng)用程序中的網(wǎng)絡(luò)具有多樣性，并從以下方面受益。

Raptor 在 NCCL 之間切換的能力，用于點(diǎn)到點(diǎn)傳輸，以將模型參數(shù)直接從學(xué)習(xí)者 GPU 傳輸?shù)酵评?GPU 。

Redis 用于異步和較小的消息，如獎(jiǎng)勵(lì)或統(tǒng)計(jì)信息。

JIT 編譯的 RPC ，用于處理高容量和低延遲請(qǐng)求，例如上載體驗(yàn)數(shù)據(jù)。

最后， Raptor 提供了支持 GPU 的數(shù)據(jù)結(jié)構(gòu)，例如重播緩沖區(qū)，該緩沖區(qū)具有多線程服務(wù)器，用于接收來(lái)自多個(gè)工作人員的經(jīng)驗(yàn)，并并行地批處理數(shù)據(jù)并將其預(yù)取到 GPU 上。

圖 4 顯示，我們的框架支持并行訓(xùn)練和數(shù)據(jù)收集，并利用 NCCL 有效地向參與者發(fā)送最新參數(shù)。

圖 4 。我們使用 Raptor 進(jìn)行解耦并行訓(xùn)練和獎(jiǎng)勵(lì)計(jì)算，以克服電路合成延遲

獎(jiǎng)勵(lì)計(jì)算

我們使用［0 ， 1］中的折衷權(quán)重 w 來(lái)組合面積和延遲目標(biāo)。我們訓(xùn)練具有不同權(quán)重的各種代理，以獲得平衡面積和延遲之間權(quán)衡的帕累托前沿設(shè)計(jì)。

RL 環(huán)境中的物理合成優(yōu)化可以生成各種解決方案，以在面積和延遲之間進(jìn)行權(quán)衡。我們應(yīng)該使用與訓(xùn)練特定代理相同的權(quán)衡權(quán)重來(lái)驅(qū)動(dòng)物理合成工具。

在循環(huán)中執(zhí)行物理合成優(yōu)化以進(jìn)行獎(jiǎng)勵(lì)計(jì)算有幾個(gè)優(yōu)點(diǎn)。

RL 代理學(xué)習(xí)直接優(yōu)化目標(biāo)技術(shù)節(jié)點(diǎn)和庫(kù)的最終電路屬性。

通過(guò)在物理合成過(guò)程中包含周圍邏輯， RL 代理可以聯(lián)合優(yōu)化目標(biāo)算術(shù)電路及其周圍邏輯的屬性。

然而，執(zhí)行物理合成是一個(gè)緩慢的過(guò)程（ 64b 加法器約 35 秒），這會(huì)大大減緩 RL 訓(xùn)練和探索。

我們將獎(jiǎng)勵(lì)計(jì)算與狀態(tài)更新解耦，因?yàn)榇碇恍枰?dāng)前前綴圖狀態(tài)來(lái)采取行動(dòng)，而不需要電路合成或之前的獎(jiǎng)勵(lì)。多虧了 Raptor ，我們可以將冗長(zhǎng)的獎(jiǎng)勵(lì)計(jì)算轉(zhuǎn)移到 CPU 工作人員池中，并行執(zhí)行物理合成，而演員代理無(wú)需等待即可在環(huán)境中穿行。

當(dāng) CPU 工作者返回獎(jiǎng)勵(lì)時(shí)，可以將轉(zhuǎn)換插入重播緩沖區(qū)。緩存合成獎(jiǎng)勵(lì)，以避免在狀態(tài)重新計(jì)數(shù)時(shí)進(jìn)行冗余計(jì)算。

后果

RL 代理純粹通過(guò)從合成電路屬性反饋的學(xué)習(xí)來(lái)學(xué)習(xí)設(shè)計(jì)電路。圖 5 顯示了使用 PrefixRL 設(shè)計(jì)的 64b 加法器電路的最新結(jié)果*， Pareto 主導(dǎo)的加法器電路來(lái)自最先進(jìn)的 EDA 工具，在面積和延遲方面。

在相同延遲下，最好的 PrefixRL 加法器的面積比 EDA 工具加法器低 25% 。這些經(jīng)過(guò)物理綜合優(yōu)化后映射到帕累托最優(yōu)加法器電路的前綴圖具有不規(guī)則結(jié)構(gòu)。

圖 5 。 PrefixRL 設(shè)計(jì)的算術(shù)電路比最先進(jìn)的 EDA 工具設(shè)計(jì)的電路更小更快。（左）電路架構(gòu)；（右）對(duì)應(yīng)的 64b 加法器電路特性圖

結(jié)論

據(jù)我們所知，這是第一種使用深度強(qiáng)化學(xué)習(xí)代理設(shè)計(jì)算術(shù)電路的方法。我們希望該方法可以成為將人工智能應(yīng)用于現(xiàn)實(shí)電路設(shè)計(jì)問題的藍(lán)圖：構(gòu)建動(dòng)作空間、狀態(tài)表示、 RL 代理模型、針對(duì)多個(gè)競(jìng)爭(zhēng)目標(biāo)進(jìn)行優(yōu)化，以及克服物理合成等緩慢的獎(jiǎng)勵(lì)計(jì)算過(guò)程。

關(guān)于作者

Rajarshi Roy 是 NVIDIA 應(yīng)用深度學(xué)習(xí)研究小組的高級(jí)研究科學(xué)家。他研究了使用深度學(xué)習(xí)、機(jī)器學(xué)習(xí)和強(qiáng)化學(xué)習(xí)改進(jìn)芯片設(shè)計(jì)、架構(gòu)和系統(tǒng)的新方法。在研究之前， Rajarshi 作為 NVIDIA GPU ASIC 團(tuán)隊(duì)的硬件工程師，對(duì)幾種 GPU 架構(gòu)的設(shè)計(jì)和驗(yàn)證做出了貢獻(xiàn)。拉賈希在斯坦福大學(xué)獲得電氣工程碩士學(xué)位。

Jonathan Raiman 是 NVIDIA 應(yīng)用深度學(xué)習(xí)研究小組的高級(jí)研究科學(xué)家，致力于大規(guī)模分布式強(qiáng)化學(xué)習(xí)和系統(tǒng)人工智能。此前，他是 OpenAI 的研究科學(xué)家，在那里他共同創(chuàng)建了 OpenAI Five ，一個(gè)超人深度強(qiáng)化學(xué)習(xí) Dota 2 機(jī)器人。在百度 SVAIL ，他共同創(chuàng)建了幾個(gè)神經(jīng)文本語(yǔ)音轉(zhuǎn)換系統(tǒng)（深度語(yǔ)音 1 、 2 和 3 ），并致力于語(yǔ)音識(shí)別（深度語(yǔ)音 2 ）和問答（全球標(biāo)準(zhǔn)化閱讀器）。他也是 DeepType 1 和 DeepType 2 （一種超人實(shí)體鏈接系統(tǒng)）的創(chuàng)建者。他正在巴黎薩克萊完成博士學(xué)位，此前在麻省理工學(xué)院獲得碩士學(xué)位。

Saad Godil 是 NVIDIA 應(yīng)用深度學(xué)習(xí)研究的主管，他領(lǐng)導(dǎo)的團(tuán)隊(duì)正在探索在我們的芯片設(shè)計(jì)和硬件項(xiàng)目中使用人工智能的新方法。在此之前，他是 NVIDIA 的 GPU 驗(yàn)證主管，在半導(dǎo)體行業(yè)有十多年的經(jīng)驗(yàn)。

審核編輯：郭婷