英伟达讲述：用AI设计芯片的优势

随着摩尔定律的放缓，在同一技术工艺节点开发其他提高芯片性能的技术变得越来越重要。我们的方法使用 AI 来设计更小、更快、更高效的电路，从而为每一代芯片提供更高的性能。

大量的算术电路阵列为 NVIDIA GPU 提供了动力，为 AI、高性能计算和计算机图形实现了前所未有的加速。因此，改进这些算术电路的设计对于提高 GPU 的性能和效率至关重要。

如果人工智能可以学会设计这些电路会怎样？在PrefixRL: Optimization of Parallel Prefix Circuits using Deep Reinforcement Learning中，我们证明了 AI 不仅可以从头开始学习设计这些电路，而且 AI 设计的电路也比最先进的电子设计的电路更小、更快。设计自动化 (EDA) 工具。最新的 NVIDIA Hopper GPU 架构拥有近 13,000 个 AI 设计的电路实例。

图 1. PrefixRL AI（左）设计的 64b 加法器电路比最先进的 EDA 工具（右）设计的电路小 25%，同时速度和功能等效。

在图 1 中，电路对应于图 5 中 PrefixRL 曲线中的 (31.4µm², 0.186ns) 点。

计算机芯片中的算术电路是使用逻辑门网络（如 NAND、NOR 和 XOR）和电线构成的。理想的电路应具有以下特点：

小： 较小的区域，以便更多电路可以安装在芯片上。

快速： 降低延迟以提高芯片的性能。

消耗更少的功率： 芯片的功耗更低。

在我们的论文中，我们关注电路面积和延迟。我们发现功耗与我们感兴趣的电路的面积密切相关。电路面积和延迟通常是相互竞争的属性，因此我们希望找到能够有效权衡这些属性的设计的帕累托前沿。简而言之，我们希望在每次延迟时都有最小面积的电路。

在 PrefixRL 中，我们专注于一类流行的算术电路，称为（并行）前缀电路（(parallel) prefix circuits.）。GPU 中的各种重要电路，例如加法器、增量器和编码器，都是前缀电路，可以在更高级别上定义为前缀图（prefix graphs）。

在这项工作中，我们特别提出了一个问题：AI 代理能否设计出好的前缀图？所有前缀图的状态空间都很大O(2^n^n)，无法使用蛮力方法进行探索。

图 2. PrefixRL 的一次迭代与 4b 电路示例

使用电路生成器将前缀图转换为带有线和逻辑门的电路。这些生成的电路然后由物理综合工具进一步优化，该工具使用物理综合优化，例如gate sizing, duplication和buffer insertion.

由于这些物理综合优化，最终的电路属性（延迟、面积和功率）不会直接从原始前缀图属性（例如电平和节点计数）转换而来。这就是为什么 AI 代理会学习设计前缀图，但会针对从前缀图生成的最终电路的属性进行优化。

我们将算术电路设计视为强化学习 (RL) 任务，在其中我们训练一个代理（agent）来优化算术电路的面积和延迟特性。对于前缀电路，我们设计了一个环境，RL 代理可以在其中添加或删除前缀图中的节点，之后会发生以下步骤：

1.前缀图被合法化以始终保持正确的前缀和计算。

2.从合法化的前缀图生成电路。

3.该电路使用物理综合工具进行物理综合优化。

4.测量电路的面积和延迟特性。

在一个环节中，RL 代理通过添加或删除节点逐步构建前缀图。在每一步，代理都会收到相应电路区域的改进和延迟作为奖励。

我们使用 Q-learning 算法来训练电路设计代理。我们对前缀图使用网格表示，其中网格中的每个元素唯一地映射到前缀节点。这种网格表示用于 Q 网络的输入和输出。输入网格中的每个元素表示节点是否存在。输出网格中的每个元素代表添加或删除节点的 Q 值。

我们为代理使用完全卷积神经网络架构，因为 Q-learning 代理的输入和输出是网格表示。代理分别预测面积和延迟属性的 Q 值，因为面积和延迟的奖励在训练期间是单独可观察的。

图 3. 某些 4b 前缀图（左）和全卷积 Q 学习代理架构（右）的表示

PrefixRL 是一项计算要求高的任务：物理模拟每个 GPU 需要 256 个 CPU，训练 64b 案例需要超过 32,000 个 GPU 小时。

我们开发了 Raptor，这是一个内部分布式强化学习平台，它特别利用 NVIDIA 硬件进行这种工业强化学习（图 4）。

Raptor 具有多项可提高可扩展性和训练速度的功能，例如作业调度、自定义网络和 GPU 感知数据结构。在 PrefixRL 的上下文中，Raptor 使得跨 CPU、GPU 和 Spot 实例的混合分配工作成为可能。

这个强化学习应用程序中的网络是多种多样的，并且受益于以下几点。

· Raptor 在 NCCL 之间切换以进行点对点传输以将模型参数直接从学习器 GPU 传输到推理 GPU 的能力。

· Redis 用于异步和较小的消息，例如奖励或统计信息。

· 一种 JIT 编译的 RPC，用于处理大容量和低延迟的请求，例如上传体验数据。

最后，Raptor 提供了 GPU 感知数据结构，例如具有多线程服务器的重播缓冲区，以接收来自多个工作人员的经验，并行批处理数据并将其预取到 GPU 上。

图 4 显示我们的框架支持并发训练和数据收集，并利用 NCCL 有效地向参与者发送最新参数。

图 4. 我们使用 Raptor 进行解耦和并行化训练和奖励计算，以克服电路合成延迟

我们使用 [0,1] 中的权衡权重w来组合区域和延迟目标。我们训练具有不同权重的各种代理以获得设计的Pareto frontier，以平衡面积和延迟之间的权衡。

RL 环境中的物理综合优化可以生成各种解决方案来权衡面积和延迟。我们应该使用与训练特定代理相同的权衡权重来驱动物理合成工具。

在奖励计算的循环中执行物理综合优化有几个优点。

· RL 代理学习直接优化目标技术节点和库的最终电路属性。

· RL 代理可以通过在物理综合过程中包含周围逻辑来联合优化目标算术电路及其周围逻辑的属性。

然而，执行物理合成是一个缓慢的过程（64b 加法器约为 35 秒），这会大大减慢 RL 训练和探索。

我们将奖励计算与状态更新解耦，因为代理只需要当前前缀图状态即可采取行动，而不需要电路合成或先前的奖励。多亏了 Raptor，我们可以将冗长的奖励计算卸载到 CPU 工作人员池中以并行执行物理合成，而参与者代理无需等待即可逐步通过环境。

当 CPU 工作人员返回奖励时，可以将转换插入到重放缓冲区中。合成奖励被缓存以避免在再次遇到状态时进行冗余计算。

RL agents纯粹通过学习合成电路属性的反馈来学习设计电路。图 5 显示了使用 PrefixRL 设计的 64b 加法器电路的最新结果*，Pareto 主导的加法器电路来自最先进的 EDA 工具的面积和延迟。

最好的 PrefixRL 加法器在相同延迟下实现的面积比 EDA 工具加法器低 25%。这些在物理综合优化后映射到Pareto optimal 加法器电路的前缀图具有不规则的结构。

图 5. PrefixRL 设计的算术电路比最先进的 EDA 工具设计的电路更小、更快。（左）电路架构；（右）相应的 64b 加法器电路特性图

据我们所知，这是第一种使用深度强化学习代理来设计算术电路的方法。我们希望这种方法可以成为将 AI 应用于现实世界电路设计问题的蓝图：构建动作空间、状态表示、RL 代理模型、针对多个竞争目标进行优化，以及克服物理合成等缓慢的奖励计算过程。