英偉達用AI設計GPU算術電路有何優(yōu)勢

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

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

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

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

最新的英偉達Hopper GPU架構中擁有近13000個AI設計的電路實例。下圖1左PrefixRL AI設計的64b加法器電路比圖1右最先進EDA工具設計的電路小25%。

電路設計概覽

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

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

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

更低功耗。

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

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

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

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

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

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

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

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

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

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

原圖為可交互版本

完全卷積Q學習智能體

研究者采用Q學習（Q-learning）算法來訓練智能體電路設計。如下圖3所示，他們將前綴圖分解成網(wǎng)格表示，其中網(wǎng)格中的每個元素唯一地映射到前綴節(jié)點。這種網(wǎng)格表示用于Q網(wǎng)絡的輸入和輸出。輸入網(wǎng)格中的每個元素表示節(jié)點是否存在。輸出網(wǎng)格中的每個元素代表添加或刪除節(jié)點的Q值。
研究者采用完全卷積神經(jīng)網(wǎng)絡架構，因為Q學習智能體的輸入和輸出都是網(wǎng)格表示。智能體分別預測面積和延遲屬性的Q值，因為面積和延遲的獎勵在訓練期間是單獨可觀察的。

圖3：4b前綴圖表示（左）以及完全卷積Q學習智能體架構（右）。 ?

Raptor進行分布式訓練

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

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

這個強化學習應用程序中的網(wǎng)絡是多種多樣的，并且受益于以下幾點：

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

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

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

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

圖4：研究者采用Raptor進行解耦并行訓練和獎勵計算，以克服電路合成延遲。

獎勵計算

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

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

RL智能體學習直接優(yōu)化目標技術節(jié)點和庫的最終電路屬性；

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

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

結(jié)果及展望

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