谷歌的 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 年,团队就发表了一篇预印本论文,介绍了谷歌的全新强化学习方法,用于设计芯片布局。
后来在 2021 年,这项工作发表在了 Nature 上,并且进行了开源。
而自从首次发表这项工作以来,谷歌内部一直在对它进行改进。今天,团队发表了 Nature 附录,详细描述了具体方法,及其对芯片设计领域的影响。同时,他们还发布了一个预训练的检查点,分享了模型权重,公布模型名称为 AlphaChip。
谷歌表示,AlphaChip 是最早用于解决现实世界工问题的强化学习方法之一。
在数小时内,它就可以生成超人或类似的芯片布局,而不需要花费数周或数月的人类。它设计的芯片布局,已经被用于世界各地的芯片中,包括数据中心和移动电话。
为了设计 TPU 布局,AlphaChip 首先在来自前几代的各种芯片模块上进行实践,例如片上和片间网络模块、内存控制器和数据传输缓冲区。这一过程被称为预训练。
然后,团队在当前的 TPU 模块上运行 AlphaChip,以生成高质量的布局。
与之前的方法不同,AlphaChip 在解决更多芯片布局任务时变得更好、更快,类似于人类专家的工作方式。
对于每一代新的 TPU,包括谷歌最新的 Trillium (第 6 代),AlphaChip 都设计了更好的芯片布局,并提供了更多的总体布局图,从而加快了设计周期,产生了更高性能的芯片。
工作原理:一边设计,一边奖励
其实,设计芯片布局并不是一项简单的任务。
一般来说,计算机芯片有许多相互连接的模块、多层电路元件组成,所有这些部件都由纤细无比的导线连接起来。
此外,还有许多复杂且相互交织的设计约束,必须同时满足。
由于设计的复杂性,60 多年来,芯片设计工程师一直在努力自动化芯片布局规划过程。
谷歌表示,AlphaChip 的研发,从 AlphaGo 和 AlphaZero 中汲取了经验。
众所周知,通过深度学习和博弈论,AlphaGo 和 AlphaZero 逐渐从 0 掌握了围棋、国际象棋和将棋的潜在规则。
AlphaChip 同样是采用了,将芯片底层规划视为一种游戏的策略。
从空白栅格开始,AlphaChip 每次放置一个电路元件,直至放置完所有元件。
然后,根据最终布局的质量,给予模型奖励。
一种全新的「基于边」的图神经网络让 AlphaChip,能够学习相互连接的芯片元件之间的关系,并在芯片之间进行泛化,让 AlphaChip 在设计的每种布局中都有所改进。
左图:动画显示 AlphaChip 在没有任何经验的情况下,将开源的 Ariane RISC-V CPU 置入。右图:动画显示 AlphaChip 在对 20 个 TPU 相关设计进行练习后,放置相同的块。
AI 大牛带队,2 页浓缩版力作
让我们从最新论文中深扒一下,AlphaChip 的整个训练过程。
值得一提的是,这项研究依旧是由 Jeff Dean 带队,所有核心要素全都浓缩在了这两页论文中。
预训练
与以往方法不同的是,AlphaChip 是基于一种「强化学习」的方法。
这意味着,当它解决了更多芯片布局问题的实例时,会变得更好、更快。
正如 Nature 论文(2021 年),以及 ISPD 2022 后续研究中所讨论的那样,这种预训练过程显著提升了 AI 的速度、可靠性、布局的性能。
顺便提一句,预训练也能培养出 LLM,像 Gemini、ChatGPT 令人印象深刻的能力。
自此前研究发表以来,谷歌便开源了一个软件库,以重现论文中描述的方法。
开发者们可以使用这个库,对各种芯片进行预训练,然后将预训练的模型应用到新的块。
基于最新的 AlphaChip 训练过程,研究人员在库中添加了预训练的 20 个 TPU 块模型检查点(checkpoint)。
显然,如果不进行任何预训练,AlphaChip 就无法从先前的经验中学习,从而规避了学习方面的问题。
训练和计算资源
随着 RL 智能体(任何 ML 模型)的投入训练,它的损失通常会逐渐减少。
最终会趋于平稳,这代表着模型对其正在执行的任务有了尽可能多的了解,对外表现就是「收敛」。
从训练到收敛,是机器学习的标准做法。如果不按照这个路径来,可能会损害模型的性能。
AlphaChip 的性能随应用的计算资源而扩展,在 ISPD 2022 论文中,谷歌团队曾进一步探讨了这一特性。
正如 Nature 论文中所描述的,在对特定块进行微调时,使用了 16 个工作单元,每个单元由 1 个 GPU 和 32 个 RL 环境组成,通过多进程处理共享 10 个 CPU。
总言之,用较少的计算资源可能会损害性能,或者需要运行相当长的时间,才能实现相同(或更差)性能。
初始布局
在运行 Nature 论文中评估方法之前,团队使用了来自物理综合的近似初始布局,以解决 hMETIS 标准单元集群大小不平衡的问题。
RL 智能体无权访问初始布局,并且不复杂放置标准单元。
尽管如此,谷歌作者还是进行了一项消融研究,排除了任何初始布局的使用,并且也没有观察到 AlphaChip 性能下降。
如下表 1 所示。
具体来说,他们跳过了单元集群重新平衡的一步,而是将 hMETIS 集群不平衡参数降低到最低设置(UBfactor = 1)。
由此,这使得 hMETIS 生成更平衡的集群。
基准
在 Nature 论文中,研究人员采用了 10nm 以下制程的 TPU 块进行实验得出的结果。
这个技术节点的大小,正是现代芯片的典型尺寸。之前许多论文报告中,采用较早的 45nm、12nm。
从物理设计角度来看,这种较老的技术节点尺寸的芯片,有着显著的不同。
比如,在 10nm 以下的芯片中,通常使用多重图案设计,这会在较低密度下导致布线拥堵的问题。
因此,对于较早的技术节点尺寸,AlphaChip 可能需要调整其奖励函数,以便更好地适应技术。
展望未来:AI 将改变整个芯片设计流程
自从 2020 年发布以来,AlphaChip 已经生成了每一代谷歌 TPU 使用的超人芯片布局。
可以说,正是因为它,才能使大规模放大基于 Transformer 架构的 AI 模型成为可能。
无论是在 Gemini 这样的 LLM,还是 Imagen 和 Veo 这样的图像和视频生成器中,TPU 都位于谷歌强大的生成式 AI 系统的核心。
另外,这些 AI 加速器也处于谷歌 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
广告声明:文内含有的对外跳转链接(包括不限于超链接、二维码、口令等形式),用于传递更多信息,节省甄选时间,结果仅供参考,IT之家所有文章均包含本声明。