AI 几小时设计芯片超越人类:谷歌 AlphaChip 登 Nature,已设计出三代旗舰 TPU

谷歌的 AlphaChip,几小时内就能设计出芯片布局,直接碾压人类专家!这种超人芯片布局,已经应用在 TPU、CPU 在内的全球硬件中。人类设计芯片的方式,已被 AI 彻底改变。

能设计芯片的 AI 黑科技来了!就在刚刚,谷歌 DeepMind 推出名为 AlphaChip 的 AI 系统。

无论是设计最先进的用于构建 AI 模型的 TPU,还是数据中心的 CPU,它在相关的众多领域,都产生了广泛影响。

在谷歌的许多款芯片设计中,它都取得了出色的效果,比如 Axion 芯片(一种基于 Arm 的通用数据中心 CPU)。

AlphaChip 设计芯片,用的是强化学习的原理。

也就是说,芯片布局设计对它来说是一种游戏,就像 AlphaGo 一样,它在游戏中,学习如何设计出最好的芯片布局。

几小时内,它就能生成超出人类水平,或是与人类专家相当的芯片布局了。

现在,它已经用于设计多代 TPU 芯片(TPU v5e、TPU v5p 和 Trillium)。而且跟人类专家相比,AlphaChip 放置的块数越来越多,线长也减少了许多。

布局五年,谷歌多代 TPU 全由 AI 设计

其实谷歌对于这个 AI,已经布局多年了。早在 2020 年,团队就发表了一篇预印本论文,介绍了谷歌的全新强化学习方法,用于设计芯片布局。

论文地址:https://arxiv.org/ pdf/2004.10746

后来在 2021 年,这项工作发表在了 Nature 上,并且进行了开源。

论文地址:https://www.nature.com/ articles / s41586-021-03544-w

而自从首次发表这项工作以来,谷歌内部一直在对它进行改进。今天,团队发表了 Nature 附录,详细描述了具体方法,及其对芯片设计领域的影响。同时,他们还发布了一个预训练的检查点,分享了模型权重,公布模型名称为 AlphaChip。

谷歌表示,AlphaChip 是最早用于解决现实世界工问题的强化学习方法之一。

在数小时内,它就可以生成超人或类似的芯片布局,而不需要花费数周或数月的人类。它设计的芯片布局,已经被用于世界各地的芯片中,包括数据中心和移动电话。

为了设计 TPU 布局,AlphaChip 首先在来自前几代的各种芯片模块上进行实践,例如片上和片间网络模块、内存控制器和数据传输缓冲区。这一过程被称为预训练。

然后,团队在当前的 TPU 模块上运行 AlphaChip,以生成高质量的布局。

与之前的方法不同,AlphaChip 在解决更多芯片布局任务时变得更好、更快,类似于人类专家的工作方式。

对于每一代新的 TPU,包括谷歌最新的 Trillium (第 6 代),AlphaChip 都设计了更好的芯片布局,并提供了更多的总体布局图,从而加快了设计周期,产生了更高性能的芯片。

条形图显示了谷歌三代 TPU 上 AlphaChip 设计的芯片块的数量,包括 v5e、v5p 和 Trillium

条形图显示,跟 TPU 物理设计团队生成的布局相比,AlphaChip 在 TPU 三代产品中的平均有线长度减少

工作原理:一边设计,一边奖励

其实,设计芯片布局并不是一项简单的任务。

一般来说,计算机芯片有许多相互连接的模块、多层电路元件组成,所有这些部件都由纤细无比的导线连接起来。

此外,还有许多复杂且相互交织的设计约束,必须同时满足。

由于设计的复杂性,60 多年来,芯片设计工程师一直在努力自动化芯片布局规划过程。

谷歌表示,AlphaChip 的研发,从 AlphaGo 和 AlphaZero 中汲取了经验。

众所周知,通过深度学习和博弈论,AlphaGo 和 AlphaZero 逐渐从 0 掌握了围棋、国际象棋和将棋的潜在规则。

AlphaChip 同样是采用了,将芯片底层规划视为一种游戏的策略。

从空白栅格开始,AlphaChip 每次放置一个电路元件,直至放置完所有元件。

然后,根据最终布局的质量,给予模型奖励。

一种全新的「基于边」的图神经网络让 AlphaChip,能够学习相互连接的芯片元件之间的关系,并在芯片之间进行泛化,让 AlphaChip 在设计的每种布局中都有所改进。

左图:动画显示 AlphaChip 在没有任何经验的情况下,将开源的 Ariane RISC-V CPU 置入。右图:动画显示 AlphaChip 在对 20 个 TPU 相关设计进行练习后,放置相同的块。

AI 大牛带队,2 页浓缩版力作

让我们从最新论文中深扒一下,AlphaChip 的整个训练过程。

值得一提的是,这项研究依旧是由 Jeff Dean 带队,所有核心要素全都浓缩在了这两页论文中。

论文地址:https://www.nature.com/ articles / s41586-024-08032-5

预训练

与以往方法不同的是,AlphaChip 是基于一种「强化学习」的方法。

这意味着,当它解决了更多芯片布局问题的实例时,会变得更好、更快。

正如 Nature 论文(2021 年),以及 ISPD 2022 后续研究中所讨论的那样,这种预训练过程显著提升了 AI 的速度、可靠性、布局的性能。

顺便提一句,预训练也能培养出 LLM,像 Gemini、ChatGPT 令人印象深刻的能力。

自此前研究发表以来,谷歌便开源了一个软件库,以重现论文中描述的方法。

开发者们可以使用这个库,对各种芯片进行预训练,然后将预训练的模型应用到新的块。

GitHub 地址:https://github.com/ google-research / circuit_training

基于最新的 AlphaChip 训练过程,研究人员在库中添加了预训练的 20 个 TPU 块模型检查点(checkpoint)。

显然,如果不进行任何预训练,AlphaChip 就无法从先前的经验中学习,从而规避了学习方面的问题。

训练和计算资源

随着 RL 智能体(任何 ML 模型)的投入训练,它的损失通常会逐渐减少。

最终会趋于平稳,这代表着模型对其正在执行的任务有了尽可能多的了解,对外表现就是「收敛」。

从训练到收敛,是机器学习的标准做法。如果不按照这个路径来,可能会损害模型的性能。

AlphaChip 的性能随应用的计算资源而扩展,在 ISPD 2022 论文中,谷歌团队曾进一步探讨了这一特性。

论文地址:https://dl.acm.org/ doi / 10.1145/3505170.3511478

正如 Nature 论文中所描述的,在对特定块进行微调时,使用了 16 个工作单元,每个单元由 1 个 GPU 和 32 个 RL 环境组成,通过多进程处理共享 10 个 CPU。

总言之,用较少的计算资源可能会损害性能,或者需要运行相当长的时间,才能实现相同(或更差)性能。

初始布局

在运行 Nature 论文中评估方法之前,团队使用了来自物理综合的近似初始布局,以解决 hMETIS 标准单元集群大小不平衡的问题。

RL 智能体无权访问初始布局,并且不复杂放置标准单元。

尽管如此,谷歌作者还是进行了一项消融研究,排除了任何初始布局的使用,并且也没有观察到 AlphaChip 性能下降。

如下表 1 所示。

具体来说,他们跳过了单元集群重新平衡的一步,而是将 hMETIS 集群不平衡参数降低到最低设置(UBfactor = 1)。

由此,这使得 hMETIS 生成更平衡的集群。

基准

在 Nature 论文中,研究人员采用了 10nm 以下制程的 TPU 块进行实验得出的结果。

这个技术节点的大小,正是现代芯片的典型尺寸。之前许多论文报告中,采用较早的 45nm、12nm。

从物理设计角度来看,这种较老的技术节点尺寸的芯片,有着显著的不同。

比如,在 10nm 以下的芯片中,通常使用多重图案设计,这会在较低密度下导致布线拥堵的问题。

来源:Pushing Multiple Patterning in Sub-10nm: Are We Ready?

因此,对于较早的技术节点尺寸,AlphaChip 可能需要调整其奖励函数,以便更好地适应技术。

展望未来:AI 将改变整个芯片设计流程

自从 2020 年发布以来,AlphaChip 已经生成了每一代谷歌 TPU 使用的超人芯片布局。

可以说,正是因为它,才能使大规模放大基于 Transformer 架构的 AI 模型成为可能。

无论是在 Gemini 这样的 LLM,还是 Imagen 和 Veo 这样的图像和视频生成器中,TPU 都位于谷歌强大的生成式 AI 系统的核心。

另外,这些 AI 加速器也处于谷歌 AI 服务的核心,外部用户可以通过谷歌云获得服务。

谷歌数据中心的一排 Cloud TPU v5p AI 加速器超算

如今,谷歌的三代旗舰 TPU 芯片,已经在世界各地的数据中心中制造、部署。

随着每一代 TPU 的发展,AlphaChip 和人类专家之间的性能差距不断扩大。

从 TPU v5e 中的 10 个 RL 放置模块和 3.2% 的布线长度减少,到 TPU v5p 中的 15 个模块和 4.5% 的减少,再到 Trillium 中的 25 个模块和 6.2% 的减少。

AlphaChip 还为数据中心 CPU(Axion)和谷歌尚未公布的其他芯片,生成了超越人类的布局设计。

而其他公司,也在谷歌研究的基础上进行了改进。

比如联发科就不仅用 AlphaChip 加速了最先进芯片的开发,还在功耗、性能和面积上对芯片做了优化。

如今,AlphaChip 仅仅是一个开始。

谷歌对未来做出了大胆畅想:AI 将实现芯片设计全流程的自动化,通过超人算法以及硬件、软件和机器学习模型的端到端协同优化,芯片设计的周期会显著加快,还会解锁性能的新领域。

谷歌表示,非常期待和社区合作,实现 AI 芯片以及芯片 AI 之间的闭环。

参考资料:

  • https://deepmind.google/discover/blog/how-alphachip-transformed-computer-chip-design/?utm_source=x&utm_medium=social&utm_campaign=&utm_content=

  • https://x.com/JeffDean/status/1839308592408834559

Published by

风君子

独自遨游何稽首 揭天掀地慰生平