隨著摩爾定律的放緩，在相同的技術(shù)工藝節(jié)點(diǎn)上開發(fā)能夠提升芯片性能的其他技術(shù)變得越來越重要。在這項(xiàng)研究中，英偉達(dá)使用深度強(qiáng)化學(xué)習(xí)方法設(shè)計(jì)尺寸更小、速度更快和更加高效的算術(shù)電路，從而為芯片提供更高的性能。

大量的算術(shù)電路陣列為英偉達(dá)GPU提供了動(dòng)力，以實(shí)現(xiàn)前所未有的AI、高性能計(jì)算和計(jì)算機(jī)圖形加速。因此，改進(jìn)這些算術(shù)電路的設(shè)計(jì)對于提升 GPU 性能和效率而言至關(guān)重要。

如果AI學(xué)習(xí)設(shè)計(jì)這些電路會(huì)怎么樣呢？在近期英偉達(dá)的論文《PrefixRL: Optimization of Parallel Prefix Circuits using Deep Reinforcement Learning》中，研究者證明了AI不僅可以從頭開始設(shè)計(jì)這些電路，而且AI設(shè)計(jì)的電路比最先進(jìn)電子設(shè)計(jì)自動(dòng)化（EDA）工具設(shè)計(jì)的電路更小、更快。

論文地址：https://arxiv.org/pdf/2205.07000.pdf

最新的英偉達(dá)Hopper GPU架構(gòu)中擁有近13000個(gè)AI設(shè)計(jì)的電路實(shí)例。下圖1左PrefixRL AI設(shè)計(jì)的64b加法器電路比圖1右最先進(jìn)EDA工具設(shè)計(jì)的電路小25%。

電路設(shè)計(jì)概覽

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

?。焊〉拿娣e，更多電路可以封裝在芯片上；

快：更低的延遲，提高芯片的性能；

更低功耗。

在英偉達(dá)的這項(xiàng)研究中，研究者關(guān)注電路面積和延遲。他們發(fā)現(xiàn)，功耗與感興趣電路的面積密切相關(guān)。電路面積和延遲往往是相互競爭的屬性，因此希望找到有效權(quán)衡這些屬性的設(shè)計(jì)的帕累托邊界。簡言之，研究者希望每次延遲時(shí)電路面積是最小的。

因此，在PrefixRL中，研究者專注于一類流行的算術(shù)電路——并行前綴電路。GPU中的各種重要電路如加速器、增量器和編碼器等都是前綴電路，它們可以在更高級(jí)別上被定為為前綴圖。

那么問題來了：AI智能體能設(shè)計(jì)出好的前綴圖嗎？所有前綴圖的狀態(tài)空間是很大的O(2^n^n)，無法使用蠻力方法進(jìn)行探索。下圖2為具有4b電路實(shí)例的PrefixRL的一次迭代。

研究者使用電路生成器將前綴圖轉(zhuǎn)換為一個(gè)帶有電線和邏輯門的電路。接下來，這些生成的電路通過一個(gè)物理綜合工具來優(yōu)化，該工具使用門尺寸、復(fù)制和緩沖器插入等物理綜合優(yōu)化。

由于這些物理綜合優(yōu)化，最終的電路屬性（延遲、面積和功率）不會(huì)直接從原始前綴圖屬性（如電平和節(jié)點(diǎn)數(shù)）轉(zhuǎn)換而來。這就是為什么AI智能體學(xué)習(xí)設(shè)計(jì)前綴圖但又要對從前綴圖中生成的最終電路的屬性進(jìn)行優(yōu)化。 研究者將算術(shù)電路設(shè)計(jì)視為一項(xiàng)強(qiáng)化學(xué)習(xí)（RL）任務(wù)，其中訓(xùn)練一個(gè)智能體優(yōu)化算術(shù)電路的面積和延遲屬性。對于前綴電路，他們設(shè)計(jì)了一個(gè)環(huán)境，其中RL智能體可以添加或刪除前綴圖中的節(jié)點(diǎn)，然后執(zhí)行如下步驟：

前綴圖被規(guī)范化以始終保持正確的前綴和計(jì)算；

從規(guī)范化的前綴圖中生成電路；

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

測量電路的面積和延遲特性。

在如下動(dòng)圖中，RL智能體通過添加或刪除節(jié)點(diǎn)來一步步地構(gòu)建前綴圖。在每一步上，該智能體得到的獎(jiǎng)勵(lì)是對應(yīng)電路面積和延遲的改進(jìn)。

原圖為可交互版本

完全卷積Q學(xué)習(xí)智能體

研究者采用Q學(xué)習(xí)（Q-learning）算法來訓(xùn)練智能體電路設(shè)計(jì)。如下圖3所示，他們將前綴圖分解成網(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)，因?yàn)镼學(xué)習(xí)智能體的輸入和輸出都是網(wǎng)格表示。智能體分別預(yù)測面積和延遲屬性的Q值，因?yàn)槊娣e和延遲的獎(jiǎng)勵(lì)在訓(xùn)練期間是單獨(dú)可觀察的。

圖3：4b前綴圖表示（左）以及完全卷積Q學(xué)習(xí)智能體架構(gòu)（右）。 ?

Raptor進(jìn)行分布式訓(xùn)練

PrefixRL需要大量計(jì)算，在物理模擬中，每個(gè)GPU需要256個(gè)CPU，而訓(xùn)練64b任務(wù)則需要超過32000個(gè)GPU小時(shí)。此次英偉達(dá)開發(fā)了一個(gè)內(nèi)分布式強(qiáng)化學(xué)習(xí)平臺(tái)Raptor，該平臺(tái)充分利用了英偉達(dá)硬件優(yōu)勢，從而可以進(jìn)行這種工業(yè)級(jí)別的強(qiáng)化學(xué)習(xí)（如下圖4）。

Raptor能夠提高訓(xùn)練模型的可擴(kuò)展性和訓(xùn)練速度，例如作業(yè)調(diào)度、自定義網(wǎng)絡(luò)和GPU感知數(shù)據(jù)結(jié)構(gòu)。在PrefixRL的上下文中，Raptor使得跨CPU、GPU和Spot實(shí)例的混合分配成為可能。

這個(gè)強(qiáng)化學(xué)習(xí)應(yīng)用程序中的網(wǎng)絡(luò)是多種多樣的，并且受益于以下幾點(diǎn)：

Raptor在NCCL之間切換以進(jìn)行點(diǎn)對點(diǎn)傳輸，從而將模型參數(shù)直接從學(xué)習(xí)器GPU傳輸?shù)酵评鞧PU；

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

對于JIT編譯的RPC，用于處理大容量和低延遲的請求，例如上傳經(jīng)驗(yàn)數(shù)據(jù)。

最后，Raptor提供了GPU感知數(shù)據(jù)結(jié)構(gòu)，例如具有多線程服務(wù)的重放緩沖器，以接收來自多個(gè)worker的經(jīng)驗(yàn)，并行批處理數(shù)據(jù)并將其預(yù)先載入到GPU上。
下圖4顯示PrefixRL框架支持并發(fā)訓(xùn)練和數(shù)據(jù)收集，并利用NCCL有效地向參與者（下圖中的actor）發(fā)送最新參數(shù)。

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

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

研究者采用權(quán)衡權(quán)重w （范圍為[0,1]）來組合區(qū)域和延遲目標(biāo)。他們訓(xùn)練具有不同權(quán)重的各種智能體以獲得帕累托邊界，從而平衡面積、延遲之間的權(quán)衡。

RL環(huán)境中的物理綜合優(yōu)化可以生成各種解決方案來權(quán)衡面積和延遲。研究者使用與訓(xùn)練特定智能體相同的權(quán)衡權(quán)重來驅(qū)動(dòng)物理綜合工具。
在獎(jiǎng)勵(lì)計(jì)算的循環(huán)中執(zhí)行物理綜合優(yōu)化具有以下優(yōu)點(diǎn)：

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

RL智能體在物理綜合過程中包含目標(biāo)算法電路的周邊邏輯，從而共同優(yōu)化目標(biāo)算法電路及其周邊邏輯的性能。

然而，進(jìn)行物理綜合是一個(gè)緩慢的過程(64b加法器~35秒)，這可能大大減慢RL的訓(xùn)練和探索。
研究者將獎(jiǎng)勵(lì)計(jì)算與狀態(tài)更新解耦，因?yàn)橹悄荏w只需要當(dāng)前的前綴圖狀態(tài)就可以采取行動(dòng)，而不需要電路合成或之前的獎(jiǎng)勵(lì)。得益于Raptor，他們可以將冗長的獎(jiǎng)勵(lì)計(jì)算轉(zhuǎn)移到CPU worker池中并行執(zhí)行物理綜合，而actor智能體無需等待就能在環(huán)境中執(zhí)行。
當(dāng)CPU worker返回獎(jiǎng)勵(lì)時(shí)，轉(zhuǎn)換就可以嵌入重放緩沖器。綜合獎(jiǎng)勵(lì)會(huì)被緩存，以避免再次遇到某個(gè)狀態(tài)時(shí)進(jìn)行冗余計(jì)算。

結(jié)果及展望

下圖5展示了使用PrefixRL設(shè)計(jì)的64b加法器電路、以及來自最先進(jìn)EDA工具的帕累托支配加法器電路的面積和延遲情況。
最好的PrefixRL加法器在相同延遲下實(shí)現(xiàn)的面積比EDA工具加法器低25%。這些在物理綜合優(yōu)化后映射到Pareto最優(yōu)加法器電路的前綴圖具有不規(guī)則的結(jié)構(gòu)。

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

據(jù)了解，這是首個(gè)使用深度強(qiáng)化學(xué)習(xí)智能體來設(shè)計(jì)算術(shù)電路的方法。英偉達(dá)構(gòu)想了一種藍(lán)圖：希望這種方法可以將AI應(yīng)用于現(xiàn)實(shí)世界電路設(shè)計(jì)問題，構(gòu)建動(dòng)作空間、狀態(tài)表示、RL 智能體模型、針對多個(gè)競爭目標(biāo)進(jìn)行優(yōu)化，以及克服緩慢的獎(jiǎng)勵(lì)計(jì)算過程。

審核編輯：黃飛

?