这篇文章最初发表在 NVIDIA 技术博客上。
随着摩尔定律的放缓，开发其他技术来提高同一技术过程节点上芯片的性能变得越来越重要。我们的方法使用人工智能设计更小、更快、更高效的电路，以在每一代芯片中提供更高的性能。
大量的算术电路阵列为 NVIDIA GPU 提供了动力，使其在人工智能、高性能计算和计算机图形学方面实现了前所未有的加速。因此，改进这些算术电路的设计对于提高 GPU 的性能和效率至关重要。
如果人工智能能够学会设计这些电路呢？在 PrefixRL ：使用深度强化学习优化并行前缀电路 中，我们证明了人工智能不仅可以从头开始学习设计这些电路，而且人工智能设计的电路也比最先进的电子设计自动化（ EDA ）工具设计的电路更小更快。最新 NVIDIA Hopper GPU 结构 拥有近 13000 个人工智能设计电路实例。
图 1 。由 PrefixRL AI （左）设计的 64b 加法器电路比由最先进的 EDA 工具（右）设计的电路小 25% ，同时速度快，功能等效
在图 1 中，电路对应于图 5 中 PrefixRL 曲线中的（ 31.4 µ m ²， 0.186ns ）点。
电路设计游戏
计算机芯片中的算术电路是使用逻辑门（如 NAND 、 NOR 和 XOR ）和导线组成的网络构建的。理想电路应具有以下特性：
- Small: 一个较低的区域，以便在一个芯片上可以容纳更多的电路。
- Fast: 降低延迟以提高芯片性能。
- 耗电更少： 芯片功耗更低。
在本文中，我们重点研究了电路面积和延迟。我们发现功耗与感兴趣电路的面积密切相关。电路面积和延迟通常是相互竞争的属性，因此我们希望找到有效权衡这些属性的设计的帕累托前沿。简单地说，我们希望在每个延迟时都有最小面积电路。
在 PrefixRL 中，我们关注一类流行的算术电路，称为（并行）前缀电路。 GPU 中的各种重要电路，如加法器、增量器和编码器，都是前缀电路，可以在更高级别上定义为前缀图。
在这项工作中，我们特别提出了一个问题：人工智能代理能否设计出良好的前缀图？所有前缀图的状态空间都很大O(2^n^n)，无法使用蛮力方法进行探索。
图 2:PrefixRL 的一次迭代与 4b 电路示例
使用电路生成器将前缀图转换为具有导线和逻辑门的电路。然后，使用物理合成优化（如栅极尺寸、复制和缓冲器插入）通过物理合成工具进一步优化这些生成的电路。
由于这些物理合成优化，最终电路特性（延迟、面积和功率）不会直接从原始前缀图特性（如电平和节点数）转换。这就是为什么人工智能代理学习设计前缀图，但优化由前缀图生成的最终电路的属性。
我们将算术电路设计作为强化学习（ RL ）任务，在该任务中，我们训练代理优化算术电路的面积和延迟特性。对于前缀电路，我们设计了一个环境，在该环境中， RL 代理可以在前缀图中添加或删除节点，然后执行以下步骤：
- 前缀图合法化，以始终保持正确的前缀和计算。
- 从合法化的前缀图生成电路。
- 使用物理合成工具对电路进行物理合成优化。
- 测量了电路的面积和延迟特性。
在一集中， RL 代理通过添加或删除节点逐步建立前缀图。在每个步骤中，代理都会收到相应电路区域的改进和延迟作为奖励。
状态和动作表示与深度强化学习模型
我们使用 Q 学习算法来训练电路设计代理。我们对前缀图使用网格表示，其中网格中的每个元素唯一地映射到前缀节点。这种网格表示法用于 Q 网络的输入和输出。输入网格中的每个元素表示节点是否存在。输出网格中的每个元素表示用于添加或删除节点的 Q 值。
我们使用完全卷积神经网络架构作为智能体的输入和输出， Q 学习智能体是网格表示。该代理单独预测面积和延迟属性的 Q 值，因为面积和延迟的回报在训练期间是可以单独观察到的。
图 3 。某些 4b 前缀图（左）和全卷积 Q 学习代理架构（右）的表示
Raptor 分布式训练
PrefixRL 是一项计算要求很高的任务：物理模拟每个 GPU 需要 256 CPU ，训练 64b 案例需要 32000 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 来组合面积和延迟目标。我们训练具有不同权重的各种代理，以获得平衡面积和延迟之间权衡的帕累托前沿设计。
RL 环境中的物理合成优化可以生成各种解决方案，以在面积和延迟之间进行权衡。我们应该使用与训练特定代理相同的权衡权重来驱动物理合成工具。
在循环中执行物理合成优化以进行奖励计算有几个优点。
- RL 代理学习直接优化目标技术节点和库的最终电路属性。
- 通过在物理合成过程中包含周围逻辑， RL 代理可以联合优化目标算术电路及其周围逻辑的属性。
然而，执行物理合成是一个缓慢的过程（ 64b 加法器约 35 秒），这会大大减缓 RL 训练和探索。
我们将奖励计算与状态更新解耦，因为代理只需要当前前缀图状态来采取行动，而不需要电路合成或之前的奖励。多亏了 Raptor ，我们可以将冗长的奖励计算转移到 CPU 工作人员池中，并行执行物理合成，而演员代理无需等待即可在环境中穿行。
当 CPU 工作者返回奖励时，可以将转换插入重播缓冲区。缓存合成奖励，以避免在状态重新计数时进行冗余计算。
后果
RL 代理纯粹通过从合成电路属性反馈的学习来学习设计电路。图 5 显示了使用 PrefixRL 设计的 64b 加法器电路的最新结果*， Pareto 主导的加法器电路来自最先进的 EDA 工具，在面积和延迟方面。
在相同延迟下，最好的 PrefixRL 加法器的面积比 EDA 工具加法器低 25% 。这些经过物理综合优化后映射到帕累托最优加法器电路的前缀图具有不规则结构。
图 5 。 PrefixRL 设计的算术电路比最先进的 EDA 工具设计的电路更小更快。（左）电路架构；（右）对应的 64b 加法器电路特性图
结论
据我们所知，这是第一种使用深度强化学习代理设计算术电路的方法。我们希望该方法可以成为将人工智能应用于现实电路设计问题的蓝图：构建动作空间、状态表示、 RL 代理模型、针对多个竞争目标进行优化，以及克服物理合成等缓慢的奖励计算过程。
有关更多信息以及与其他方法的比较，请参阅 PrefixRL ：使用深度强化学习优化并行前缀电路 （ preprint ）。